sochdb_kernel/

atomic_claim.rs

// SPDX-License-Identifier: AGPL-3.0-or-later
// SochDB - LLM-Optimized Embedded Database
// Copyright (C) 2026 Sushanth Reddy Vanagala (https://github.com/sushanthpy)
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Atomic Claim Protocol for Queue Operations
//!
//! This module provides linearizable "claim + delete" pop semantics for queue
//! operations, ensuring no double-delivery under concurrent access.
//!
//! ## Problem: Double-Delivery Under Concurrency
//!
//! Without atomic claims, a naive "scan → delete" pattern can fail:
//!
//! ```text
//! Worker A: scan() → finds task T
//! Worker B: scan() → finds task T  (same task!)
//! Worker A: delete(T) → success
//! Worker B: delete(T) → fails or double-processes
//! ```
//!
//! ## Solution: CAS-Based Claim Protocol
//!
//! The claim is the linearization point. Only one worker can successfully
//! create the claim key, establishing ownership:
//!
//! ```text
//! Worker A: scan() → finds task T
//! Worker A: CAS(claim/T, absent → A) → SUCCESS
//! Worker B: scan() → finds task T
//! Worker B: CAS(claim/T, absent → B) → FAIL (key exists)
//! Worker B: retry with next candidate
//! ```
//!
//! ## Lease-Based Crash Recovery
//!
//! Claims have expiry times. If a worker crashes:
//!
//! 1. Claim expires after `lease_duration`
//! 2. Next worker's scan finds task with expired claim
//! 3. New claim can overwrite expired claim
//! 4. Task is reprocessed (at-least-once delivery)
//!
//! ## Integration with SochDB's MVCC/SSI
//!
//! The claim protocol works with SochDB's transaction model:
//!
//! - Claims use MVCC versioning for conflict detection
//! - The claim key insert is the durability boundary
//! - SSI's dangerous-structure detection catches anomalies
//! - WAL + fsync ensures claim survives crashes

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use parking_lot::{Mutex, RwLock};

// ============================================================================
// ClaimResult - Result of a Claim Attempt
// ============================================================================

/// Result of attempting to claim a task
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ClaimResult {
    /// Successfully claimed the task
    Success {
        /// The claim token for subsequent operations
        claim_token: ClaimToken,
    },
    /// Task was already claimed by another worker
    AlreadyClaimed {
        /// Who holds the claim
        owner: String,
        /// When the claim expires
        expires_at: u64,
    },
    /// Claim expired and was taken over
    TookOver {
        /// Previous owner whose claim expired
        previous_owner: String,
        /// New claim token
        claim_token: ClaimToken,
    },
    /// Task not found
    NotFound,
    /// Internal error
    Error(String),
}

// ============================================================================
// ClaimToken - Proof of Ownership
// ============================================================================

/// A token proving ownership of a claimed task
///
/// This token must be presented for subsequent operations (ack, nack, extend).
/// It prevents workers from operating on tasks they don't own.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ClaimToken {
    /// Queue containing the task
    pub queue_id: String,
    /// Task being claimed
    pub task_id: String,
    /// Owner identity
    pub owner: String,
    /// Unique claim instance (to detect stale tokens)
    pub instance: u64,
    /// When the claim was created (epoch millis)
    pub created_at: u64,
    /// When the claim expires (epoch millis)
    pub expires_at: u64,
}

impl ClaimToken {
    /// Check if this token is still valid
    pub fn is_valid(&self, now_millis: u64) -> bool {
        now_millis < self.expires_at
    }

    /// Time remaining on the lease
    pub fn remaining_ms(&self, now_millis: u64) -> u64 {
        self.expires_at.saturating_sub(now_millis)
    }
}

// ============================================================================
// ClaimEntry - Internal Claim State
// ============================================================================

/// Internal state of a claim
#[derive(Debug, Clone)]
struct ClaimEntry {
    /// Owner identity
    owner: String,
    /// Unique instance ID (increments on each claim)
    instance: u64,
    /// When claimed (epoch millis)
    claimed_at: u64,
    /// When claim expires (epoch millis)
    expires_at: u64,
    /// Number of times this task has been claimed
    claim_count: u32,
}

impl ClaimEntry {
    fn is_expired(&self, now_millis: u64) -> bool {
        now_millis >= self.expires_at
    }

    fn to_token(&self, queue_id: &str, task_id: &str) -> ClaimToken {
        ClaimToken {
            queue_id: queue_id.to_string(),
            task_id: task_id.to_string(),
            owner: self.owner.clone(),
            instance: self.instance,
            created_at: self.claimed_at,
            expires_at: self.expires_at,
        }
    }
}

// ============================================================================
// AtomicClaimManager - The Core Claim Coordination Layer
// ============================================================================

/// Atomic claim manager for queue task ownership
///
/// This provides the CAS-based claim protocol that ensures linearizable
/// task ownership under concurrent access.
///
/// ## Thread Safety
///
/// All operations are thread-safe. The manager uses fine-grained locking
/// to minimize contention:
/// - Per-queue locks for claim operations
/// - Read-write locks for statistics
///
/// ## Durability
///
/// In production, claims should be persisted to storage with WAL durability.
/// This in-memory implementation is for reference and testing.
pub struct AtomicClaimManager {
    /// Claims by queue_id -> (task_id -> ClaimEntry)
    claims: RwLock<HashMap<String, HashMap<String, ClaimEntry>>>,
    /// Instance counter for unique claim IDs
    instance_counter: AtomicU64,
    /// Statistics
    stats: RwLock<ClaimStats>,
    /// Mutex for claim operations (ensures CAS semantics)
    claim_locks: RwLock<HashMap<String, std::sync::Arc<Mutex<()>>>>,
}

/// Statistics for claim operations
#[derive(Debug, Clone, Default)]
pub struct ClaimStats {
    /// Total claim attempts
    pub attempts: u64,
    /// Successful claims
    pub successes: u64,
    /// Failed due to contention
    pub contentions: u64,
    /// Takeovers of expired claims
    pub takeovers: u64,
    /// Claims released via ack
    pub acks: u64,
    /// Claims released via nack
    pub nacks: u64,
    /// Claims expired
    pub expirations: u64,
}

impl Default for AtomicClaimManager {
    fn default() -> Self {
        Self::new()
    }
}

impl AtomicClaimManager {
    /// Create a new claim manager
    pub fn new() -> Self {
        Self {
            claims: RwLock::new(HashMap::new()),
            instance_counter: AtomicU64::new(1),
            stats: RwLock::new(ClaimStats::default()),
            claim_locks: RwLock::new(HashMap::new()),
        }
    }

    /// Get or create a lock for a specific claim key
    fn get_claim_lock(&self, queue_id: &str, task_id: &str) -> std::sync::Arc<Mutex<()>> {
        let key = format!("{}:{}", queue_id, task_id);

        // Fast path: check if lock exists
        {
            let locks = self.claim_locks.read();
            if let Some(lock) = locks.get(&key) {
                return lock.clone();
            }
        }

        // Slow path: create lock
        let mut locks = self.claim_locks.write();
        locks
            .entry(key)
            .or_insert_with(|| std::sync::Arc::new(Mutex::new(())))
            .clone()
    }

    /// Attempt to claim a task
    ///
    /// This is the atomic CAS operation that establishes ownership.
    ///
    /// ## Semantics
    ///
    /// - If task is unclaimed: creates claim, returns Success
    /// - If task is claimed by other worker with valid lease: returns AlreadyClaimed
    /// - If task is claimed but lease expired: creates new claim, returns TookOver
    ///
    /// ## Complexity
    ///
    /// O(1) hash lookups + lock acquisition
    pub fn claim(
        &self,
        queue_id: &str,
        task_id: &str,
        owner: &str,
        lease_duration_ms: u64,
    ) -> ClaimResult {
        let now = current_time_millis();

        // Get per-claim lock to ensure CAS semantics
        let lock = self.get_claim_lock(queue_id, task_id);
        let _guard = lock.lock();

        // Update stats
        self.stats.write().attempts += 1;

        let mut claims = self.claims.write();
        let queue_claims = claims
            .entry(queue_id.to_string())
            .or_insert_with(HashMap::new);

        // Check existing claim
        if let Some(existing) = queue_claims.get(task_id) {
            if existing.owner == owner {
                // Same owner re-claiming (extend)
                let instance = self.instance_counter.fetch_add(1, AtomicOrdering::SeqCst);
                let new_entry = ClaimEntry {
                    owner: owner.to_string(),
                    instance,
                    claimed_at: now,
                    expires_at: now + lease_duration_ms,
                    claim_count: existing.claim_count + 1,
                };
                let token = new_entry.to_token(queue_id, task_id);
                queue_claims.insert(task_id.to_string(), new_entry);

                self.stats.write().successes += 1;
                return ClaimResult::Success { claim_token: token };
            }

            if !existing.is_expired(now) {
                // Valid claim by another worker
                self.stats.write().contentions += 1;
                return ClaimResult::AlreadyClaimed {
                    owner: existing.owner.clone(),
                    expires_at: existing.expires_at,
                };
            }

            // Expired claim - take over
            let previous_owner = existing.owner.clone();
            let instance = self.instance_counter.fetch_add(1, AtomicOrdering::SeqCst);
            let new_entry = ClaimEntry {
                owner: owner.to_string(),
                instance,
                claimed_at: now,
                expires_at: now + lease_duration_ms,
                claim_count: existing.claim_count + 1,
            };
            let token = new_entry.to_token(queue_id, task_id);
            queue_claims.insert(task_id.to_string(), new_entry);

            self.stats.write().takeovers += 1;
            return ClaimResult::TookOver {
                previous_owner,
                claim_token: token,
            };
        }

        // No existing claim - create new
        let instance = self.instance_counter.fetch_add(1, AtomicOrdering::SeqCst);
        let entry = ClaimEntry {
            owner: owner.to_string(),
            instance,
            claimed_at: now,
            expires_at: now + lease_duration_ms,
            claim_count: 1,
        };
        let token = entry.to_token(queue_id, task_id);
        queue_claims.insert(task_id.to_string(), entry);

        self.stats.write().successes += 1;
        ClaimResult::Success { claim_token: token }
    }

    /// Release a claim (acknowledge successful processing)
    ///
    /// The claim token must be valid and owned by the caller.
    pub fn release(&self, token: &ClaimToken) -> Result<(), String> {
        let _now = current_time_millis();

        let lock = self.get_claim_lock(&token.queue_id, &token.task_id);
        let _guard = lock.lock();

        let mut claims = self.claims.write();

        if let Some(queue_claims) = claims.get_mut(&token.queue_id) {
            if let Some(existing) = queue_claims.get(&token.task_id) {
                // Verify ownership
                if existing.instance != token.instance {
                    return Err("Stale claim token".to_string());
                }
                if existing.owner != token.owner {
                    return Err("Not claim owner".to_string());
                }

                queue_claims.remove(&token.task_id);
                self.stats.write().acks += 1;
                return Ok(());
            }
        }

        Err("Claim not found".to_string())
    }

    /// Extend a claim's lease duration
    ///
    /// Useful when processing takes longer than expected.
    pub fn extend(
        &self,
        queue_id: &str,
        token: &ClaimToken,
        additional_ms: u64,
    ) -> Result<ClaimToken, String> {
        let _now = current_time_millis();

        let lock = self.get_claim_lock(queue_id, &token.task_id);
        let _guard = lock.lock();

        let mut claims = self.claims.write();

        if let Some(queue_claims) = claims.get_mut(queue_id) {
            if let Some(existing) = queue_claims.get_mut(&token.task_id) {
                // Verify ownership
                if existing.instance != token.instance {
                    return Err("Stale claim token".to_string());
                }
                if existing.owner != token.owner {
                    return Err("Not claim owner".to_string());
                }

                // Extend the lease
                existing.expires_at += additional_ms;

                return Ok(existing.to_token(queue_id, &token.task_id));
            }
        }

        Err("Claim not found".to_string())
    }

    /// Check if a task is currently claimed
    pub fn is_claimed(&self, queue_id: &str, task_id: &str) -> Option<(String, u64)> {
        let now = current_time_millis();

        let claims = self.claims.read();

        if let Some(queue_claims) = claims.get(queue_id) {
            if let Some(entry) = queue_claims.get(task_id) {
                if !entry.is_expired(now) {
                    return Some((entry.owner.clone(), entry.expires_at));
                }
            }
        }

        None
    }

    /// Get the current claim token for a task (if owned by the given worker)
    pub fn get_token(&self, queue_id: &str, task_id: &str, owner: &str) -> Option<ClaimToken> {
        let now = current_time_millis();

        let claims = self.claims.read();

        if let Some(queue_claims) = claims.get(queue_id) {
            if let Some(entry) = queue_claims.get(task_id) {
                if !entry.is_expired(now) && entry.owner == owner {
                    return Some(entry.to_token(queue_id, task_id));
                }
            }
        }

        None
    }

    /// Clean up expired claims
    ///
    /// This should be called periodically (e.g., every few seconds).
    /// Returns the number of claims cleaned up.
    pub fn cleanup_expired(&self) -> usize {
        let now = current_time_millis();
        let mut cleaned = 0;

        let mut claims = self.claims.write();

        for queue_claims in claims.values_mut() {
            queue_claims.retain(|_, entry| {
                if entry.is_expired(now) {
                    cleaned += 1;
                    false
                } else {
                    true
                }
            });
        }

        if cleaned > 0 {
            self.stats.write().expirations += cleaned as u64;
        }

        cleaned
    }

    /// Get statistics
    pub fn stats(&self) -> ClaimStats {
        self.stats.read().clone()
    }

    /// Get number of active claims for a queue
    pub fn active_claims(&self, queue_id: &str) -> usize {
        let now = current_time_millis();

        self.claims
            .read()
            .get(queue_id)
            .map(|q| q.values().filter(|e| !e.is_expired(now)).count())
            .unwrap_or(0)
    }

    /// Get all active claims for a queue (for monitoring)
    pub fn list_claims(&self, queue_id: &str) -> Vec<ClaimToken> {
        let now = current_time_millis();

        self.claims
            .read()
            .get(queue_id)
            .map(|q| {
                q.iter()
                    .filter(|(_, e)| !e.is_expired(now))
                    .map(|(task_id, e)| e.to_token(queue_id, task_id))
                    .collect()
            })
            .unwrap_or_default()
    }
}

// ============================================================================
// CompareAndSwap Trait - For Storage Integration
// ============================================================================

/// Compare-and-swap trait for storage backends
///
/// This trait abstracts the CAS operation for different storage implementations.
/// SochDB's storage layer should implement this for durable claims.
pub trait CompareAndSwap {
    /// Type of error returned
    type Error: std::fmt::Debug;

    /// Insert a key-value pair only if the key doesn't exist
    ///
    /// Returns Ok(true) if inserted, Ok(false) if key exists, Err on failure.
    fn insert_if_absent(&self, key: &[u8], value: &[u8]) -> Result<bool, Self::Error>;

    /// Update a value only if the current value matches expected
    ///
    /// Returns Ok(true) if updated, Ok(false) if mismatch, Err on failure.
    fn compare_and_set(
        &self,
        key: &[u8],
        expected: &[u8],
        new_value: &[u8],
    ) -> Result<bool, Self::Error>;

    /// Delete a key only if the current value matches expected
    ///
    /// Returns Ok(true) if deleted, Ok(false) if mismatch, Err on failure.
    fn delete_if_match(&self, key: &[u8], expected: &[u8]) -> Result<bool, Self::Error>;
}

// ============================================================================
// LeaseManager - Higher-Level Lease Coordination
// ============================================================================

/// Configuration for lease management
#[derive(Debug, Clone)]
pub struct LeaseConfig {
    /// Default lease duration
    pub default_lease_ms: u64,
    /// Minimum lease duration
    pub min_lease_ms: u64,
    /// Maximum lease duration
    pub max_lease_ms: u64,
    /// How often to run cleanup (ms)
    pub cleanup_interval_ms: u64,
    /// Maximum extensions per task
    pub max_extensions: u32,
}

impl Default for LeaseConfig {
    fn default() -> Self {
        Self {
            default_lease_ms: 30_000,   // 30 seconds
            min_lease_ms: 1_000,        // 1 second
            max_lease_ms: 3_600_000,    // 1 hour
            cleanup_interval_ms: 5_000, // 5 seconds
            max_extensions: 10,
        }
    }
}

/// Higher-level lease manager with periodic cleanup
pub struct LeaseManager {
    /// Underlying claim manager
    claim_manager: AtomicClaimManager,
    /// Configuration
    config: LeaseConfig,
    /// Last cleanup time
    last_cleanup: RwLock<Instant>,
    /// Extension counts per task
    extension_counts: RwLock<HashMap<String, u32>>,
}

impl LeaseManager {
    /// Create a new lease manager
    pub fn new(config: LeaseConfig) -> Self {
        Self {
            claim_manager: AtomicClaimManager::new(),
            config,
            last_cleanup: RwLock::new(Instant::now()),
            extension_counts: RwLock::new(HashMap::new()),
        }
    }

    /// Acquire a lease on a task
    pub fn acquire(
        &self,
        queue_id: &str,
        task_id: &str,
        owner: &str,
        lease_ms: Option<u64>,
    ) -> ClaimResult {
        self.maybe_cleanup();

        let lease_duration = lease_ms
            .unwrap_or(self.config.default_lease_ms)
            .clamp(self.config.min_lease_ms, self.config.max_lease_ms);

        self.claim_manager
            .claim(queue_id, task_id, owner, lease_duration)
    }

    /// Release a lease
    pub fn release(&self, queue_id: &str, token: &ClaimToken) -> Result<(), String> {
        // Clear extension count
        {
            let key = format!("{}:{}", queue_id, token.task_id);
            self.extension_counts.write().remove(&key);
        }

        // Delegate to claim manager (which now correctly uses token.queue_id)
        self.claim_manager.release(token)
    }

    /// Extend a lease
    pub fn extend(
        &self,
        queue_id: &str,
        token: &ClaimToken,
        additional_ms: u64,
    ) -> Result<ClaimToken, String> {
        let key = format!("{}:{}", queue_id, token.task_id);

        // Check extension limit
        {
            let counts = self.extension_counts.read();
            if let Some(&count) = counts.get(&key) {
                if count >= self.config.max_extensions {
                    return Err(format!(
                        "Maximum extensions ({}) reached",
                        self.config.max_extensions
                    ));
                }
            }
        }

        // Clamp additional time
        let additional = additional_ms.clamp(self.config.min_lease_ms, self.config.max_lease_ms);

        let result = self.claim_manager.extend(queue_id, token, additional)?;

        // Increment extension count
        {
            let mut counts = self.extension_counts.write();
            *counts.entry(key).or_insert(0) += 1;
        }

        Ok(result)
    }

    /// Get claim manager statistics
    pub fn stats(&self) -> ClaimStats {
        self.claim_manager.stats()
    }

    /// Force cleanup of expired leases
    pub fn cleanup(&self) -> usize {
        *self.last_cleanup.write() = Instant::now();
        self.claim_manager.cleanup_expired()
    }

    /// Check if cleanup should run and run it if needed
    fn maybe_cleanup(&self) {
        let should_cleanup = {
            let last = self.last_cleanup.read();
            last.elapsed() > Duration::from_millis(self.config.cleanup_interval_ms)
        };

        if should_cleanup {
            self.cleanup();
        }
    }
}

// ============================================================================
// Helper Functions
// ============================================================================

/// Get current time in milliseconds since epoch
fn current_time_millis() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as u64
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use std::thread;

    #[test]
    fn test_claim_success() {
        let manager = AtomicClaimManager::new();

        match manager.claim("queue1", "task1", "worker1", 30_000) {
            ClaimResult::Success { claim_token } => {
                assert_eq!(claim_token.task_id, "task1");
                assert_eq!(claim_token.owner, "worker1");
            }
            _ => panic!("Expected success"),
        }
    }

    #[test]
    fn test_claim_contention() {
        let manager = AtomicClaimManager::new();

        // First claim succeeds
        let result1 = manager.claim("queue1", "task1", "worker1", 30_000);
        assert!(matches!(result1, ClaimResult::Success { .. }));

        // Second claim fails
        let result2 = manager.claim("queue1", "task1", "worker2", 30_000);
        match result2 {
            ClaimResult::AlreadyClaimed { owner, .. } => {
                assert_eq!(owner, "worker1");
            }
            _ => panic!("Expected AlreadyClaimed"),
        }
    }

    #[test]
    fn test_claim_takeover() {
        let manager = AtomicClaimManager::new();

        // Create claim with very short lease
        let result1 = manager.claim("queue1", "task1", "worker1", 1);
        assert!(matches!(result1, ClaimResult::Success { .. }));

        // Wait for expiration
        thread::sleep(Duration::from_millis(10));

        // New worker can take over
        let result2 = manager.claim("queue1", "task1", "worker2", 30_000);
        match result2 {
            ClaimResult::TookOver { previous_owner, .. } => {
                assert_eq!(previous_owner, "worker1");
            }
            _ => panic!("Expected TookOver, got {:?}", result2),
        }
    }

    #[test]
    fn test_concurrent_claims() {
        let manager = Arc::new(AtomicClaimManager::new());
        let successes = Arc::new(std::sync::atomic::AtomicUsize::new(0));

        let mut handles = vec![];

        for i in 0..10 {
            let mgr = manager.clone();
            let succ = successes.clone();

            handles.push(thread::spawn(move || {
                match mgr.claim("queue1", "task1", &format!("worker{}", i), 30_000) {
                    ClaimResult::Success { .. } => {
                        succ.fetch_add(1, AtomicOrdering::SeqCst);
                    }
                    _ => {}
                }
            }));
        }

        for h in handles {
            h.join().unwrap();
        }

        // Only one worker should succeed
        assert_eq!(successes.load(AtomicOrdering::SeqCst), 1);
    }

    #[test]
    fn test_claim_release_wrong_queue() {
        let manager = AtomicClaimManager::new();

        // Claim task in queue1
        let token = match manager.claim("queue1", "task1", "worker1", 30_000) {
            ClaimResult::Success { claim_token } => claim_token,
            _ => panic!("Expected success"),
        };

        // Token should have correct queue_id
        assert_eq!(token.queue_id, "queue1");

        // Verify claimed
        assert!(manager.is_claimed("queue1", "task1").is_some());

        // Release using the correct token should work
        manager.release(&token).unwrap();

        // After release, task should not be claimed
        assert!(manager.is_claimed("queue1", "task1").is_none());
    }

    #[test]
    fn test_multiple_queue_isolation() {
        let manager = AtomicClaimManager::new();

        // Claim same task_id in different queues
        let token1 = match manager.claim("queue1", "task1", "worker1", 30_000) {
            ClaimResult::Success { claim_token } => claim_token,
            _ => panic!("Expected success"),
        };

        let token2 = match manager.claim("queue2", "task1", "worker1", 30_000) {
            ClaimResult::Success { claim_token } => claim_token,
            _ => panic!("Expected success"),
        };

        // Tokens should have distinct queue_ids
        assert_eq!(token1.queue_id, "queue1");
        assert_eq!(token2.queue_id, "queue2");

        // Both tasks should be claimed
        assert!(manager.is_claimed("queue1", "task1").is_some());
        assert!(manager.is_claimed("queue2", "task1").is_some());

        // Release from queue1 should not affect queue2
        manager.release(&token1).unwrap();
        assert!(manager.is_claimed("queue1", "task1").is_none());
        assert!(manager.is_claimed("queue2", "task1").is_some());

        // Release from queue2
        manager.release(&token2).unwrap();
        assert!(manager.is_claimed("queue2", "task1").is_none());
    }

    #[test]
    fn test_lease_manager_extension_limit() {
        let config = LeaseConfig {
            max_extensions: 2,
            default_lease_ms: 100,
            min_lease_ms: 10,
            max_lease_ms: 1000,
            cleanup_interval_ms: 10000,
        };

        let manager = LeaseManager::new(config);

        let token = match manager.acquire("queue1", "task1", "worker1", None) {
            ClaimResult::Success { claim_token } => claim_token,
            _ => panic!("Expected success"),
        };

        // First extension OK
        let token = manager.extend("queue1", &token, 100).unwrap();

        // Second extension OK
        let token = manager.extend("queue1", &token, 100).unwrap();

        // Third extension should fail
        let result = manager.extend("queue1", &token, 100);
        assert!(result.is_err());
    }

    #[test]
    fn test_cleanup_expired() {
        let manager = AtomicClaimManager::new();

        // Create claims with very short leases
        manager.claim("queue1", "task1", "worker1", 1);
        manager.claim("queue1", "task2", "worker1", 1);
        manager.claim("queue1", "task3", "worker1", 100_000); // Long lease

        thread::sleep(Duration::from_millis(10));

        let cleaned = manager.cleanup_expired();
        assert_eq!(cleaned, 2); // task1 and task2 expired

        // task3 should still be claimed
        assert!(manager.is_claimed("queue1", "task3").is_some());
    }

    #[test]
    fn test_stats_tracking() {
        let manager = AtomicClaimManager::new();

        // Success
        manager.claim("queue1", "task1", "worker1", 30_000);

        // Contention
        manager.claim("queue1", "task1", "worker2", 30_000);

        let stats = manager.stats();
        assert_eq!(stats.attempts, 2);
        assert_eq!(stats.successes, 1);
        assert_eq!(stats.contentions, 1);
    }
}