Fix lock (#510)

* Refactor: reimplement lock Signed-off-by: junxiang Mu <1948535941@qq.com> * Fix: fix test case failed Signed-off-by: junxiang Mu <1948535941@qq.com> * Improve: lock pref Signed-off-by: junxiang Mu <1948535941@qq.com> * fix(lock): Fix resource cleanup issue when batch lock acquisition fails Ensure that the locks already acquired are properly released when batch lock acquisition fails to avoid memory leaks Improve the lock protection mechanism to prevent double release issues Add complete Apache license declarations to all files Signed-off-by: junxiang Mu <1948535941@qq.com> --------- Signed-off-by: junxiang Mu <1948535941@qq.com>
2026-08-22 12:26:37 +00:00 · 2025-09-11 12:10:35 +08:00
parent 971e74281c
commit d4beb1cc0b
24 changed files with 4281 additions and 1349 deletions
@@ -0,0 +1,505 @@
+// Copyright 2024 RustFS Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+use tokio::sync::RwLock;
+use tokio::time::{Instant, interval};
+
+use crate::fast_lock::{
+    guard::FastLockGuard,
+    metrics::GlobalMetrics,
+    shard::LockShard,
+    types::{BatchLockRequest, BatchLockResult, LockConfig, LockResult, ObjectLockRequest},
+};
+
+/// High-performance object lock manager
+#[derive(Debug)]
+pub struct FastObjectLockManager {
+    shards: Vec<Arc<LockShard>>,
+    shard_mask: usize,
+    config: LockConfig,
+    metrics: Arc<GlobalMetrics>,
+    cleanup_handle: RwLock<Option<tokio::task::JoinHandle<()>>>,
+}
+
+impl FastObjectLockManager {
+    /// Create new lock manager with default config
+    pub fn new() -> Self {
+        Self::with_config(LockConfig::default())
+    }
+
+    /// Create new lock manager with custom config
+    pub fn with_config(config: LockConfig) -> Self {
+        let shard_count = config.shard_count;
+        assert!(shard_count.is_power_of_two(), "Shard count must be power of 2");
+
+        let shards: Vec<Arc<LockShard>> = (0..shard_count).map(|i| Arc::new(LockShard::new(i))).collect();
+
+        let metrics = Arc::new(GlobalMetrics::new(shard_count));
+
+        let manager = Self {
+            shards,
+            shard_mask: shard_count - 1,
+            config,
+            metrics,
+            cleanup_handle: RwLock::new(None),
+        };
+
+        // Start background cleanup task
+        manager.start_cleanup_task();
+        manager
+    }
+
+    /// Acquire object lock
+    pub async fn acquire_lock(&self, request: ObjectLockRequest) -> Result<FastLockGuard, LockResult> {
+        let shard = self.get_shard(&request.key);
+        match shard.acquire_lock(&request).await {
+            Ok(()) => Ok(FastLockGuard::new(request.key, request.mode, request.owner, shard.clone())),
+            Err(err) => Err(err),
+        }
+    }
+
+    /// Acquire shared (read) lock
+    pub async fn acquire_read_lock(
+        &self,
+        bucket: impl Into<Arc<str>>,
+        object: impl Into<Arc<str>>,
+        owner: impl Into<Arc<str>>,
+    ) -> Result<FastLockGuard, LockResult> {
+        let request = ObjectLockRequest::new_read(bucket, object, owner);
+        self.acquire_lock(request).await
+    }
+
+    /// Acquire shared (read) lock for specific version
+    pub async fn acquire_read_lock_versioned(
+        &self,
+        bucket: impl Into<Arc<str>>,
+        object: impl Into<Arc<str>>,
+        version: impl Into<Arc<str>>,
+        owner: impl Into<Arc<str>>,
+    ) -> Result<FastLockGuard, LockResult> {
+        let request = ObjectLockRequest::new_read(bucket, object, owner).with_version(version);
+        self.acquire_lock(request).await
+    }
+
+    /// Acquire exclusive (write) lock
+    pub async fn acquire_write_lock(
+        &self,
+        bucket: impl Into<Arc<str>>,
+        object: impl Into<Arc<str>>,
+        owner: impl Into<Arc<str>>,
+    ) -> Result<FastLockGuard, LockResult> {
+        let request = ObjectLockRequest::new_write(bucket, object, owner);
+        self.acquire_lock(request).await
+    }
+
+    /// Acquire exclusive (write) lock for specific version
+    pub async fn acquire_write_lock_versioned(
+        &self,
+        bucket: impl Into<Arc<str>>,
+        object: impl Into<Arc<str>>,
+        version: impl Into<Arc<str>>,
+        owner: impl Into<Arc<str>>,
+    ) -> Result<FastLockGuard, LockResult> {
+        let request = ObjectLockRequest::new_write(bucket, object, owner).with_version(version);
+        self.acquire_lock(request).await
+    }
+
+    /// Acquire multiple locks atomically - optimized version
+    pub async fn acquire_locks_batch(&self, batch_request: BatchLockRequest) -> BatchLockResult {
+        // Pre-sort requests by (shard_id, key) to avoid deadlocks
+        let mut sorted_requests = batch_request.requests;
+        sorted_requests.sort_unstable_by(|a, b| {
+            let shard_a = a.key.shard_index(self.shard_mask);
+            let shard_b = b.key.shard_index(self.shard_mask);
+            shard_a.cmp(&shard_b).then_with(|| a.key.cmp(&b.key))
+        });
+
+        // Try to use stack-allocated vectors for small batches, fallback to heap if needed
+        let shard_groups = self.group_requests_by_shard(sorted_requests);
+
+        // Choose strategy based on request type
+        if batch_request.all_or_nothing {
+            self.acquire_locks_two_phase_commit(&shard_groups).await
+        } else {
+            self.acquire_locks_best_effort(&shard_groups).await
+        }
+    }
+
+    /// Group requests by shard with proper fallback handling
+    fn group_requests_by_shard(
+        &self,
+        requests: Vec<ObjectLockRequest>,
+    ) -> std::collections::HashMap<usize, Vec<ObjectLockRequest>> {
+        let mut shard_groups = std::collections::HashMap::new();
+
+        for request in requests {
+            let shard_id = request.key.shard_index(self.shard_mask);
+            shard_groups.entry(shard_id).or_insert_with(Vec::new).push(request);
+        }
+
+        shard_groups
+    }
+
+    /// Best effort acquisition (allows partial success)
+    async fn acquire_locks_best_effort(
+        &self,
+        shard_groups: &std::collections::HashMap<usize, Vec<ObjectLockRequest>>,
+    ) -> BatchLockResult {
+        let mut all_successful = Vec::new();
+        let mut all_failed = Vec::new();
+
+        for (&shard_id, requests) in shard_groups {
+            let shard = &self.shards[shard_id];
+
+            // Try fast path first for each request
+            for request in requests {
+                if shard.try_fast_path_only(request) {
+                    all_successful.push(request.key.clone());
+                } 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)),
+                    }
+                }
+            }
+        }
+
+        let all_acquired = all_failed.is_empty();
+        BatchLockResult {
+            successful_locks: all_successful,
+            failed_locks: all_failed,
+            all_acquired,
+        }
+    }
+
+    /// Two-phase commit for atomic acquisition
+    async fn acquire_locks_two_phase_commit(
+        &self,
+        shard_groups: &std::collections::HashMap<usize, Vec<ObjectLockRequest>>,
+    ) -> BatchLockResult {
+        // Phase 1: Try to acquire all locks
+        let mut acquired_locks = Vec::new();
+        let mut failed_locks = Vec::new();
+
+        'outer: for (&shard_id, requests) in shard_groups {
+            let shard = &self.shards[shard_id];
+
+            for request in requests {
+                match shard.acquire_lock(request).await {
+                    Ok(()) => {
+                        acquired_locks.push((request.key.clone(), request.mode, request.owner.clone()));
+                    }
+                    Err(err) => {
+                        failed_locks.push((request.key.clone(), err));
+                        break 'outer; // Stop on first failure
+                    }
+                }
+            }
+        }
+
+        // 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);
+            }
+
+            return BatchLockResult {
+                successful_locks: Vec::new(),
+                failed_locks,
+                all_acquired: false,
+            };
+        }
+
+        // All successful
+        BatchLockResult {
+            successful_locks: acquired_locks.into_iter().map(|(key, _, _)| key).collect(),
+            failed_locks: Vec::new(),
+            all_acquired: true,
+        }
+    }
+
+    /// Get lock information for monitoring
+    pub fn get_lock_info(&self, key: &crate::fast_lock::types::ObjectKey) -> Option<crate::fast_lock::types::ObjectLockInfo> {
+        let shard = self.get_shard(key);
+        shard.get_lock_info(key)
+    }
+
+    /// Get aggregated metrics
+    pub fn get_metrics(&self) -> crate::fast_lock::metrics::AggregatedMetrics {
+        let shard_metrics: Vec<_> = self.shards.iter().map(|shard| shard.metrics().snapshot()).collect();
+
+        self.metrics.aggregate_shard_metrics(&shard_metrics)
+    }
+
+    /// Get total number of active locks across all shards
+    pub fn total_lock_count(&self) -> usize {
+        self.shards.iter().map(|shard| shard.lock_count()).sum()
+    }
+
+    /// Get pool statistics from all shards
+    pub fn get_pool_stats(&self) -> Vec<(u64, u64, u64, usize)> {
+        self.shards.iter().map(|shard| shard.pool_stats()).collect()
+    }
+
+    /// Force cleanup of expired locks using adaptive strategy
+    pub async fn cleanup_expired(&self) -> usize {
+        let mut total_cleaned = 0;
+
+        for shard in &self.shards {
+            total_cleaned += shard.adaptive_cleanup();
+        }
+
+        self.metrics.record_cleanup_run(total_cleaned);
+        total_cleaned
+    }
+
+    /// Force cleanup with traditional strategy (for compatibility)
+    pub async fn cleanup_expired_traditional(&self) -> usize {
+        let max_idle_millis = self.config.max_idle_time.as_millis() as u64;
+        let mut total_cleaned = 0;
+
+        for shard in &self.shards {
+            total_cleaned += shard.cleanup_expired_millis(max_idle_millis);
+        }
+
+        self.metrics.record_cleanup_run(total_cleaned);
+        total_cleaned
+    }
+
+    /// Shutdown the lock manager and cleanup resources
+    pub async fn shutdown(&self) {
+        if let Some(handle) = self.cleanup_handle.write().await.take() {
+            handle.abort();
+        }
+
+        // Final cleanup
+        self.cleanup_expired().await;
+    }
+
+    /// Get shard for object key
+    fn get_shard(&self, key: &crate::fast_lock::types::ObjectKey) -> &Arc<LockShard> {
+        let index = key.shard_index(self.shard_mask);
+        &self.shards[index]
+    }
+
+    /// Start background cleanup task
+    fn start_cleanup_task(&self) {
+        let shards = self.shards.clone();
+        let metrics = self.metrics.clone();
+        let cleanup_interval = self.config.cleanup_interval;
+        let _max_idle_time = self.config.max_idle_time;
+
+        let handle = tokio::spawn(async move {
+            let mut interval = interval(cleanup_interval);
+
+            loop {
+                interval.tick().await;
+
+                let start = Instant::now();
+                let mut total_cleaned = 0;
+
+                // Use adaptive cleanup for better performance
+                for shard in &shards {
+                    total_cleaned += shard.adaptive_cleanup();
+                }
+
+                if total_cleaned > 0 {
+                    metrics.record_cleanup_run(total_cleaned);
+                    tracing::debug!("Cleanup completed: {} objects cleaned in {:?}", total_cleaned, start.elapsed());
+                }
+            }
+        });
+
+        // Store handle for shutdown
+        if let Ok(mut cleanup_handle) = self.cleanup_handle.try_write() {
+            *cleanup_handle = Some(handle);
+        }
+    }
+}
+
+impl Default for FastObjectLockManager {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+// Implement Drop to ensure cleanup
+impl Drop for FastObjectLockManager {
+    fn drop(&mut self) {
+        // Note: We can't use async in Drop, so we just abort the cleanup task
+        if let Ok(handle_guard) = self.cleanup_handle.try_read() {
+            if let Some(handle) = handle_guard.as_ref() {
+                handle.abort();
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tokio::time::Duration;
+
+    #[tokio::test]
+    async fn test_manager_basic_operations() {
+        let manager = FastObjectLockManager::new();
+
+        // Test read lock
+        let read_guard = manager
+            .acquire_read_lock("bucket", "object", "owner1")
+            .await
+            .expect("Failed to acquire read lock");
+
+        // Should be able to acquire another read lock
+        let read_guard2 = manager
+            .acquire_read_lock("bucket", "object", "owner2")
+            .await
+            .expect("Failed to acquire second read lock");
+
+        drop(read_guard);
+        drop(read_guard2);
+
+        // Test write lock
+        let write_guard = manager
+            .acquire_write_lock("bucket", "object", "owner1")
+            .await
+            .expect("Failed to acquire write lock");
+
+        drop(write_guard);
+    }
+
+    #[tokio::test]
+    async fn test_manager_contention() {
+        let manager = Arc::new(FastObjectLockManager::new());
+
+        // Acquire write lock
+        let write_guard = manager
+            .acquire_write_lock("bucket", "object", "owner1")
+            .await
+            .expect("Failed to acquire write lock");
+
+        // Try to acquire read lock (should timeout)
+        let manager_clone = manager.clone();
+        let read_result =
+            tokio::time::timeout(Duration::from_millis(100), manager_clone.acquire_read_lock("bucket", "object", "owner2")).await;
+
+        assert!(read_result.is_err()); // Should timeout
+
+        drop(write_guard);
+
+        // Now read lock should succeed
+        let read_guard = manager
+            .acquire_read_lock("bucket", "object", "owner2")
+            .await
+            .expect("Failed to acquire read lock after write lock released");
+
+        drop(read_guard);
+    }
+
+    #[tokio::test]
+    async fn test_versioned_locks() {
+        let manager = FastObjectLockManager::new();
+
+        // Acquire lock on version v1
+        let v1_guard = manager
+            .acquire_write_lock_versioned("bucket", "object", "v1", "owner1")
+            .await
+            .expect("Failed to acquire v1 lock");
+
+        // Should be able to acquire lock on version v2 simultaneously
+        let v2_guard = manager
+            .acquire_write_lock_versioned("bucket", "object", "v2", "owner2")
+            .await
+            .expect("Failed to acquire v2 lock");
+
+        drop(v1_guard);
+        drop(v2_guard);
+    }
+
+    #[tokio::test]
+    async fn test_batch_operations() {
+        let manager = FastObjectLockManager::new();
+
+        let batch = BatchLockRequest::new("owner")
+            .add_read_lock("bucket", "obj1")
+            .add_write_lock("bucket", "obj2")
+            .with_all_or_nothing(true);
+
+        let result = manager.acquire_locks_batch(batch).await;
+        assert!(result.all_acquired);
+        assert_eq!(result.successful_locks.len(), 2);
+        assert!(result.failed_locks.is_empty());
+    }
+
+    #[tokio::test]
+    async fn test_metrics() {
+        let manager = FastObjectLockManager::new();
+
+        // Perform some operations
+        let _guard1 = manager.acquire_read_lock("bucket", "obj1", "owner").await.unwrap();
+        let _guard2 = manager.acquire_write_lock("bucket", "obj2", "owner").await.unwrap();
+
+        let metrics = manager.get_metrics();
+        assert!(metrics.shard_metrics.total_acquisitions() > 0);
+        assert!(metrics.shard_metrics.fast_path_rate() > 0.0);
+    }
+
+    #[tokio::test]
+    async fn test_cleanup() {
+        let config = LockConfig {
+            max_idle_time: Duration::from_secs(1), // Use 1 second for easier testing
+            ..Default::default()
+        };
+        let manager = FastObjectLockManager::with_config(config);
+
+        // Acquire and release some locks
+        {
+            let _guard = manager.acquire_read_lock("bucket", "obj1", "owner1").await.unwrap();
+            let _guard2 = manager.acquire_read_lock("bucket", "obj2", "owner2").await.unwrap();
+        } // Locks are released here
+
+        // Check lock count before cleanup
+        let count_before = manager.total_lock_count();
+        assert!(count_before >= 2, "Should have at least 2 locks before cleanup");
+
+        // Wait for idle timeout
+        tokio::time::sleep(Duration::from_secs(2)).await;
+
+        // Force cleanup with traditional method to ensure cleanup for testing
+        let cleaned = manager.cleanup_expired_traditional().await;
+
+        let count_after = manager.total_lock_count();
+
+        // The test should pass if cleanup works at all
+        assert!(
+            cleaned > 0 || count_after < count_before,
+            "Cleanup should either clean locks or they should be cleaned by other means"
+        );
+    }
+}