feat(lock): enhance lock management with timeout and ownership tracking (#589)

- Add lock timeout support and track acquisition time in lock state
- Improve lock conflict handling with detailed error messages
- Optimize lock reuse when already held by same owner
- Refactor lock state to store owner info and timeout duration
- Update all lock operations to handle new state structure

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

crates/lock/src/fast_lock/object_pool.rs

@@ -111,6 +111,9 @@ impl ObjectLockState {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::fast_lock::state::{ExclusiveOwnerInfo, SharedOwnerEntry};
+    use std::sync::Arc;
+    use std::time::{Duration, SystemTime};
 
     #[test]
     fn test_object_pool() {
@@ -142,8 +145,17 @@ mod tests {
         let mut state = ObjectLockState::new();
 
         // Modify state
-        *state.current_owner.write() = Some("test_owner".into());
-        state.shared_owners.write().push("shared_owner".into());
+        *state.current_owner.write() = Some(ExclusiveOwnerInfo {
+            owner: Arc::from("test_owner"),
+            acquired_at: SystemTime::now(),
+            lock_timeout: Duration::from_secs(30),
+        });
+        state.shared_owners.write().push(SharedOwnerEntry {
+            owner: Arc::from("shared_owner"),
+            count: 1,
+            acquired_at: SystemTime::now(),
+            lock_timeout: Duration::from_secs(30),
+        });
 
         // Reset
         state.reset_for_reuse();

crates/lock/src/fast_lock/shard.rs

@@ -88,8 +88,8 @@ impl LockShard {
 
                 // Try atomic acquisition
                 let success = match request.mode {
-                    LockMode::Shared => state.try_acquire_shared_fast(&request.owner),
-                    LockMode::Exclusive => state.try_acquire_exclusive_fast(&request.owner),
+                    LockMode::Shared => state.try_acquire_shared_fast(&request.owner, request.lock_timeout),
+                    LockMode::Exclusive => state.try_acquire_exclusive_fast(&request.owner, request.lock_timeout),
                 };
 
                 if success {
@@ -108,14 +108,14 @@ impl LockShard {
                 let state = state.clone();
                 drop(objects);
 
-                if state.try_acquire_exclusive_fast(&request.owner) {
+                if state.try_acquire_exclusive_fast(&request.owner, request.lock_timeout) {
                     return Some(state);
                 }
             } else {
                 // Create new state from pool and acquire immediately
                 let state_box = self.object_pool.acquire();
                 let state = Arc::new(*state_box);
-                if state.try_acquire_exclusive_fast(&request.owner) {
+                if state.try_acquire_exclusive_fast(&request.owner, request.lock_timeout) {
                     objects.insert(request.key.clone(), state.clone());
                     return Some(state);
                 }
@@ -151,8 +151,8 @@ impl LockShard {
 
             // Try acquisition again
             let success = match request.mode {
-                LockMode::Shared => state.try_acquire_shared_fast(&request.owner),
-                LockMode::Exclusive => state.try_acquire_exclusive_fast(&request.owner),
+                LockMode::Shared => state.try_acquire_shared_fast(&request.owner, request.lock_timeout),
+                LockMode::Exclusive => state.try_acquire_exclusive_fast(&request.owner, request.lock_timeout),
             };
 
             if success {
@@ -443,22 +443,24 @@ impl LockShard {
         let objects = self.objects.read();
         if let Some(state) = objects.get(key) {
             if let Some(mode) = state.current_mode() {
-                let owner = match mode {
+                let (owner, acquired_at, lock_timeout) = match mode {
                     LockMode::Exclusive => {
                         let current_owner = state.current_owner.read();
-                        current_owner.clone()?
+                        let info = current_owner.clone()?;
+                        (info.owner, info.acquired_at, info.lock_timeout)
                     }
                     LockMode::Shared => {
                         let shared_owners = state.shared_owners.read();
-                        shared_owners.first()?.clone()
+                        let entry = shared_owners.first()?.clone();
+                        (entry.owner, entry.acquired_at, entry.lock_timeout)
                     }
                 };
 
                 let priority = *state.priority.read();
 
-                // Estimate acquisition time (approximate)
-                let acquired_at = SystemTime::now() - Duration::from_secs(60);
-                let expires_at = acquired_at + Duration::from_secs(300);
+                let expires_at = acquired_at
+                    .checked_add(lock_timeout)
+                    .unwrap_or_else(|| acquired_at + crate::fast_lock::DEFAULT_LOCK_TIMEOUT);
 
                 return Some(crate::fast_lock::types::ObjectLockInfo {
                     key: key.clone(),

crates/lock/src/fast_lock/state.rs

@@ -308,13 +308,28 @@ pub struct ObjectLockState {
 
     // Third cache line: Less frequently accessed data
     /// Current owner of exclusive lock (if any)
-    pub current_owner: parking_lot::RwLock<Option<Arc<str>>>,
+    pub current_owner: parking_lot::RwLock<Option<ExclusiveOwnerInfo>>,
     /// Shared owners - optimized for small number of readers
-    pub shared_owners: parking_lot::RwLock<smallvec::SmallVec<[Arc<str>; 4]>>,
+    pub shared_owners: parking_lot::RwLock<smallvec::SmallVec<[SharedOwnerEntry; 4]>>,
     /// Lock priority for conflict resolution
     pub priority: parking_lot::RwLock<LockPriority>,
 }
 
+#[derive(Clone, Debug)]
+pub struct ExclusiveOwnerInfo {
+    pub owner: Arc<str>,
+    pub acquired_at: SystemTime,
+    pub lock_timeout: Duration,
+}
+
+#[derive(Clone, Debug)]
+pub struct SharedOwnerEntry {
+    pub owner: Arc<str>,
+    pub count: u32,
+    pub acquired_at: SystemTime,
+    pub lock_timeout: Duration,
+}
+
 impl Default for ObjectLockState {
     fn default() -> Self {
         Self::new()
@@ -335,60 +350,87 @@ impl ObjectLockState {
     }
 
     /// Try fast path shared lock acquisition
-    pub fn try_acquire_shared_fast(&self, owner: &Arc<str>) -> bool {
-        if self.atomic_state.try_acquire_shared() {
-            self.atomic_state.update_access_time();
-            let mut shared = self.shared_owners.write();
-            if !shared.contains(owner) {
-                shared.push(owner.clone());
-            }
-            true
-        } else {
-            false
+    pub fn try_acquire_shared_fast(&self, owner: &Arc<str>, lock_timeout: Duration) -> bool {
+        if !self.atomic_state.try_acquire_shared() {
+            return false;
         }
+
+        self.atomic_state.update_access_time();
+        let mut shared = self.shared_owners.write();
+        if let Some(entry) = shared.iter_mut().find(|entry| entry.owner.as_ref() == owner.as_ref()) {
+            entry.count = entry.count.saturating_add(1);
+            entry.acquired_at = SystemTime::now();
+            entry.lock_timeout = lock_timeout;
+        } else {
+            shared.push(SharedOwnerEntry {
+                owner: owner.clone(),
+                count: 1,
+                acquired_at: SystemTime::now(),
+                lock_timeout,
+            });
+        }
+        true
     }
 
     /// Try fast path exclusive lock acquisition
-    pub fn try_acquire_exclusive_fast(&self, owner: &Arc<str>) -> bool {
-        if self.atomic_state.try_acquire_exclusive() {
-            self.atomic_state.update_access_time();
-            let mut current = self.current_owner.write();
-            *current = Some(owner.clone());
-            true
-        } else {
-            false
+    pub fn try_acquire_exclusive_fast(&self, owner: &Arc<str>, lock_timeout: Duration) -> bool {
+        if !self.atomic_state.try_acquire_exclusive() {
+            return false;
         }
+
+        self.atomic_state.update_access_time();
+        let mut current = self.current_owner.write();
+        *current = Some(ExclusiveOwnerInfo {
+            owner: owner.clone(),
+            acquired_at: SystemTime::now(),
+            lock_timeout,
+        });
+        true
     }
 
     /// Release shared lock
     pub fn release_shared(&self, owner: &Arc<str>) -> bool {
         let mut shared = self.shared_owners.write();
-        if let Some(pos) = shared.iter().position(|x| x.as_ref() == owner.as_ref()) {
-            shared.remove(pos);
+        if let Some(pos) = shared.iter().position(|entry| entry.owner.as_ref() == owner.as_ref()) {
+            let original_entry = shared[pos].clone();
+            let removed_entry = if shared[pos].count > 1 {
+                shared[pos].count -= 1;
+                None
+            } else {
+                Some(shared.remove(pos))
+            };
             if self.atomic_state.release_shared() {
-                // Notify waiting writers if no more readers
                 if shared.is_empty() {
                     drop(shared);
                     self.optimized_notify.notify_writer();
                 }
                 true
             } else {
-                // Inconsistency detected - atomic state shows no shared lock but owner was found
                 tracing::warn!(
-                    "Atomic state inconsistency during shared lock release: owner={}, remaining_owners={}",
+                    "Atomic state inconsistency during shared lock release: owner={}, remaining_entries={}",
                     owner,
                     shared.len()
                 );
-                // Re-add owner to maintain consistency
-                shared.push(owner.clone());
+                // Re-add owner entry to maintain consistency when release failed
+                match removed_entry {
+                    Some(entry) => {
+                        shared.push(entry);
+                    }
+                    None => {
+                        if let Some(existing) = shared.iter_mut().find(|existing| existing.owner.as_ref() == owner.as_ref()) {
+                            existing.count = existing.count.saturating_add(1);
+                        } else {
+                            shared.push(original_entry);
+                        }
+                    }
+                }
                 false
             }
         } else {
-            // Owner not found in shared owners list
             tracing::debug!(
-                "Shared lock release failed - owner not found: owner={}, current_owners={:?}",
+                "Shared lock release failed - owner not found: owner={}, current_entries={:?}",
                 owner,
-                shared.iter().map(|s| s.as_ref()).collect::<Vec<_>>()
+                shared.iter().map(|s| s.owner.as_ref()).collect::<Vec<_>>()
             );
             false
         }
@@ -397,7 +439,7 @@ impl ObjectLockState {
     /// Release exclusive lock
     pub fn release_exclusive(&self, owner: &Arc<str>) -> bool {
         let mut current = self.current_owner.write();
-        if current.as_ref() == Some(owner) {
+        if current.as_ref().is_some_and(|info| info.owner.as_ref() == owner.as_ref()) {
             if self.atomic_state.release_exclusive() {
                 *current = None;
                 drop(current);
@@ -426,7 +468,7 @@ impl ObjectLockState {
             tracing::debug!(
                 "Exclusive lock release failed - owner mismatch: expected_owner={}, actual_owner={:?}",
                 owner,
-                current.as_ref().map(|s| s.as_ref())
+                current.as_ref().map(|s| s.owner.as_ref())
             );
             false
         }
@@ -483,16 +525,18 @@ mod tests {
         let owner2 = Arc::from("owner2");
 
         // Test shared locks
-        assert!(state.try_acquire_shared_fast(&owner1));
-        assert!(state.try_acquire_shared_fast(&owner2));
-        assert!(!state.try_acquire_exclusive_fast(&owner1));
+        let timeout = Duration::from_secs(30);
+
+        assert!(state.try_acquire_shared_fast(&owner1, timeout));
+        assert!(state.try_acquire_shared_fast(&owner2, timeout));
+        assert!(!state.try_acquire_exclusive_fast(&owner1, timeout));
 
         assert!(state.release_shared(&owner1));
         assert!(state.release_shared(&owner2));
 
         // Test exclusive lock
-        assert!(state.try_acquire_exclusive_fast(&owner1));
-        assert!(!state.try_acquire_shared_fast(&owner2));
+        assert!(state.try_acquire_exclusive_fast(&owner1, timeout));
+        assert!(!state.try_acquire_shared_fast(&owner2, timeout));
         assert!(state.release_exclusive(&owner1));
     }
 }