Improve lock (#596)

* improve lock

Signed-off-by: Mu junxiang <1948535941@qq.com>

* feat(tests): add wait_for_object_absence helper and improve lifecycle test reliability

Signed-off-by: Mu junxiang <1948535941@qq.com>

* chore: remove dirty docs

Signed-off-by: Mu junxiang <1948535941@qq.com>

---------

Signed-off-by: Mu junxiang <1948535941@qq.com>

crates/lock/src/fast_lock/disabled_manager.rs

@@ -98,12 +98,18 @@ impl DisabledLockManager {
 
     /// Always succeeds - all locks acquired
     pub async fn acquire_locks_batch(&self, batch_request: BatchLockRequest) -> BatchLockResult {
-        let successful_locks: Vec<ObjectKey> = batch_request.requests.into_iter().map(|req| req.key).collect();
+        let successful_locks: Vec<ObjectKey> = batch_request.requests.iter().map(|req| req.key.clone()).collect();
+        let guards = batch_request
+            .requests
+            .into_iter()
+            .map(|req| FastLockGuard::new_disabled(req.key, req.mode, req.owner))
+            .collect();
 
         BatchLockResult {
             successful_locks,
             failed_locks: Vec::new(),
             all_acquired: true,
+            guards,
         }
     }
 

crates/lock/src/fast_lock/manager.rs

@@ -217,20 +217,33 @@ impl FastObjectLockManager {
     ) -> BatchLockResult {
         let mut all_successful = Vec::new();
         let mut all_failed = Vec::new();
+        let mut guards = Vec::new();
 
         for (&shard_id, requests) in shard_groups {
-            let shard = &self.shards[shard_id];
+            let shard = self.shards[shard_id].clone();
 
-            // Try fast path first for each request
             for request in requests {
-                if shard.try_fast_path_only(request) {
-                    all_successful.push(request.key.clone());
+                let key = request.key.clone();
+                let owner = request.owner.clone();
+                let mode = request.mode;
+
+                let acquired = if shard.try_fast_path_only(request) {
+                    true
                 } else {
-                    // Fallback to slow path
                     match shard.acquire_lock(request).await {
-                        Ok(()) => all_successful.push(request.key.clone()),
-                        Err(err) => all_failed.push((request.key.clone(), err)),
+                        Ok(()) => true,
+                        Err(err) => {
+                            all_failed.push((key.clone(), err));
+                            false
+                        }
                     }
+                };
+
+                if acquired {
+                    let guard = FastLockGuard::new(key.clone(), mode, owner.clone(), shard.clone());
+                    shard.register_guard(guard.guard_id());
+                    all_successful.push(key);
+                    guards.push(guard);
                 }
             }
         }
@@ -240,6 +253,7 @@ impl FastObjectLockManager {
             successful_locks: all_successful,
             failed_locks: all_failed,
             all_acquired,
+            guards,
         }
     }
 
@@ -249,16 +263,18 @@ impl FastObjectLockManager {
         shard_groups: &std::collections::HashMap<usize, Vec<ObjectLockRequest>>,
     ) -> BatchLockResult {
         // Phase 1: Try to acquire all locks
-        let mut acquired_locks = Vec::new();
+        let mut acquired_guards = Vec::new();
         let mut failed_locks = Vec::new();
 
         'outer: for (&shard_id, requests) in shard_groups {
-            let shard = &self.shards[shard_id];
+            let shard = self.shards[shard_id].clone();
 
             for request in requests {
                 match shard.acquire_lock(request).await {
                     Ok(()) => {
-                        acquired_locks.push((request.key.clone(), request.mode, request.owner.clone()));
+                        let guard = FastLockGuard::new(request.key.clone(), request.mode, request.owner.clone(), shard.clone());
+                        shard.register_guard(guard.guard_id());
+                        acquired_guards.push(guard);
                     }
                     Err(err) => {
                         failed_locks.push((request.key.clone(), err));
@@ -270,35 +286,22 @@ impl FastObjectLockManager {
 
         // Phase 2: If any failed, release all acquired locks with error tracking
         if !failed_locks.is_empty() {
-            let mut cleanup_failures = 0;
-            for (key, mode, owner) in acquired_locks {
-                let shard = self.get_shard(&key);
-                if !shard.release_lock(&key, &owner, mode) {
-                    cleanup_failures += 1;
-                    tracing::warn!(
-                        "Failed to release lock during batch cleanup: bucket={}, object={}",
-                        key.bucket,
-                        key.object
-                    );
-                }
-            }
-
-            if cleanup_failures > 0 {
-                tracing::error!("Batch lock cleanup had {} failures", cleanup_failures);
-            }
-
+            // Drop guards to release any acquired locks.
+            drop(acquired_guards);
             return BatchLockResult {
                 successful_locks: Vec::new(),
                 failed_locks,
                 all_acquired: false,
+                guards: Vec::new(),
             };
         }
 
-        // All successful
+        let successful_locks = acquired_guards.iter().map(|guard| guard.key().clone()).collect();
         BatchLockResult {
-            successful_locks: acquired_locks.into_iter().map(|(key, _, _)| key).collect(),
+            successful_locks,
             failed_locks: Vec::new(),
             all_acquired: true,
+            guards: acquired_guards,
         }
     }
 

crates/lock/src/fast_lock/types.rs

@@ -19,6 +19,8 @@ use std::hash::{Hash, Hasher};
 use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 
+use crate::fast_lock::guard::FastLockGuard;
+
 /// Object key for version-aware locking
 #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct ObjectKey {
@@ -340,6 +342,7 @@ pub struct BatchLockResult {
     pub successful_locks: Vec<ObjectKey>,
     pub failed_locks: Vec<(ObjectKey, LockResult)>,
     pub all_acquired: bool,
+    pub guards: Vec<FastLockGuard>,
 }
 
 #[cfg(test)]