fileloft_store_memory/
locker.rs1use 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#[derive(Clone)]
19pub struct MemoryLocker {
20 held: HeldSet,
21 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
69pub 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 if let Ok(mut held) = self.held.try_lock() {
90 held.remove(&self.id);
91 }
92 }
95 }
96}