Skip to main content

fileloft_store_memory/
locker.rs

1use std::collections::HashSet;
2use std::sync::Arc;
3use std::time::Duration;
4
5use fileloft_core::{
6    error::TusError,
7    info::UploadId,
8    lock::{SendLock, SendLocker},
9};
10use tokio::sync::Mutex;
11
12type HeldSet = Arc<Mutex<HashSet<String>>>;
13
14/// An in-memory locker using a simple spin-wait with small sleep intervals.
15///
16/// This locker is process-local; for multi-process deployments use a
17/// distributed lock (e.g. Redis) or the file-based `FileLocker`.
18#[derive(Clone)]
19pub struct MemoryLocker {
20    held: HeldSet,
21    /// How long to wait before giving up with `LockTimeout`.
22    pub timeout: Duration,
23}
24
25impl MemoryLocker {
26    pub fn new() -> Self {
27        Self {
28            held: Arc::new(Mutex::new(HashSet::new())),
29            timeout: Duration::from_secs(20),
30        }
31    }
32
33    pub fn with_timeout(mut self, timeout: Duration) -> Self {
34        self.timeout = timeout;
35        self
36    }
37}
38
39impl Default for MemoryLocker {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45impl SendLocker for MemoryLocker {
46    type LockType = MemoryLock;
47
48    async fn acquire(&self, id: &UploadId) -> Result<MemoryLock, TusError> {
49        let deadline = tokio::time::Instant::now() + self.timeout;
50        loop {
51            {
52                let mut held = self.held.lock().await;
53                if held.insert(id.as_str().to_string()) {
54                    return Ok(MemoryLock {
55                        id: id.as_str().to_string(),
56                        held: Arc::clone(&self.held),
57                        released: false,
58                    });
59                }
60            }
61            if tokio::time::Instant::now() >= deadline {
62                return Err(TusError::LockTimeout(id.to_string()));
63            }
64            tokio::time::sleep(Duration::from_millis(10)).await;
65        }
66    }
67}
68
69/// A held in-memory lock.
70/// Automatically released on drop; also releasable explicitly via `release()`.
71pub struct MemoryLock {
72    id: String,
73    held: HeldSet,
74    released: bool,
75}
76
77impl SendLock for MemoryLock {
78    async fn release(mut self) -> Result<(), TusError> {
79        self.held.lock().await.remove(&self.id);
80        self.released = true;
81        Ok(())
82    }
83}
84
85impl Drop for MemoryLock {
86    fn drop(&mut self) {
87        if !self.released {
88            // Best-effort synchronous release: try_lock avoids blocking in Drop.
89            if let Ok(mut held) = self.held.try_lock() {
90                held.remove(&self.id);
91            }
92            // If try_lock fails, the lock will remain held until the Mutex is
93            // next acquired — acceptable in test contexts.
94        }
95    }
96}