distributed_lock_mongo/
lock.rs

1use std::time::Duration;
2
3use distributed_lock_core::{
4    error::{LockError, LockResult},
5    traits::DistributedLock,
6};
7use mongodb::{
8    Collection, Database,
9    bson::{DateTime, doc},
10    options::ReturnDocument,
11};
12use tokio::sync::watch;
13use uuid::Uuid;
14
15use crate::{
16    document::MongoLockDocument, handle::MongoLockHandle, options::MongoDistributedLockOptions,
17};
18
19pub struct MongoDistributedLock {
20    name: String,
21    collection: Collection<MongoLockDocument>,
22    options: MongoDistributedLockOptions,
23}
24
25impl MongoDistributedLock {
26    pub fn new(
27        name: String,
28        database: Database,
29        collection_name: Option<String>,
30        options: Option<MongoDistributedLockOptions>,
31    ) -> Self {
32        let collection_name = collection_name.as_deref().unwrap_or("DistributedLocks");
33        let collection = database.collection(collection_name);
34        Self {
35            name,
36            collection,
37            options: options.unwrap_or_default(),
38        }
39    }
40
41    async fn try_acquire_internal(&self) -> LockResult<Option<MongoLockHandle>> {
42        let lock_id = Uuid::new_v4().to_string();
43        let expiry_ms = self.options.expiry.as_millis() as i64;
44
45        // MongoDB Pipeline construction
46        // expired := ifNull(expiresAt, epoch) <= $$NOW
47        let epoch = DateTime::from_millis(0);
48
49        let expired_or_missing = doc! {
50            "$lte": [
51                { "$ifNull": ["$expiresAt", epoch] },
52                "$$NOW"
53            ]
54        };
55
56        let new_expires_at = doc! {
57            "$dateAdd": {
58                "startDate": "$$NOW",
59                "unit": "millisecond",
60                "amount": expiry_ms
61            }
62        };
63
64        let set_stage = doc! {
65            "$set": {
66                "lockId": {
67                    "$cond": [&expired_or_missing, &lock_id, "$lockId"]
68                },
69                "expiresAt": {
70                    "$cond": [&expired_or_missing, &new_expires_at, "$expiresAt"]
71                },
72                "acquiredAt": {
73                    "$cond": [&expired_or_missing, "$$NOW", "$acquiredAt"]
74                }
75            }
76        };
77
78        let pipeline = vec![set_stage];
79
80        let filter = doc! { "_id": &self.name };
81
82        // Use builder pattern for options in mongodb v3
83        let result = self
84            .collection
85            .find_one_and_update(filter, pipeline)
86            .upsert(true)
87            .return_document(ReturnDocument::After)
88            .await
89            .map_err(|e| LockError::Connection(Box::new(e)))?;
90
91        if let Some(doc) = result
92            && doc.lock_id == lock_id
93        {
94            let (_tx, rx) = watch::channel(false);
95            // Success!
96            return Ok(Some(MongoLockHandle {
97                collection: self.collection.clone(),
98                name: self.name.clone(),
99                lock_id,
100                lost_token: rx,
101            }));
102        }
103
104        Ok(None)
105    }
106}
107
108impl DistributedLock for MongoDistributedLock {
109    type Handle = MongoLockHandle;
110
111    fn name(&self) -> &str {
112        &self.name
113    }
114
115    async fn acquire(&self, timeout: Option<Duration>) -> LockResult<Self::Handle> {
116        let start = std::time::Instant::now();
117        let timeout = timeout.unwrap_or(Duration::from_secs(u64::MAX)); // Infinite-ish
118
119        loop {
120            if let Some(handle) = self.try_acquire_internal().await? {
121                return Ok(handle);
122            }
123
124            if start.elapsed() >= timeout {
125                return Err(LockError::Timeout(timeout));
126            }
127
128            // Simple busy wait with randomization could be better, but fixed for now
129            // C# uses randomized backoff
130            let sleep_time = self.options.min_busy_wait_sleep_time;
131            // We could implement exponential backoff here up to max_busy_wait
132            tokio::time::sleep(sleep_time).await;
133        }
134    }
135
136    async fn try_acquire(&self) -> LockResult<Option<Self::Handle>> {
137        self.try_acquire_internal().await
138    }
139}
distributed_lock_mongo/lock.rs

distributed_lock_mongo/
lock.rs
