teaql_runtime/context/
locking.rs

1use std::collections::HashMap;
2use std::sync::atomic::{AtomicU64, Ordering};
3use std::sync::{Arc, Condvar, Mutex, OnceLock};
4use std::time::{Duration, Instant};
5
6use super::UserContext;
7
8#[derive(Clone, Copy)]
9struct LocalLockEntry {
10    owner: u64,
11    expires_at: Option<Instant>,
12}
13
14#[derive(Default)]
15struct ProcessLocalLocks {
16    entries: Mutex<HashMap<String, LocalLockEntry>>,
17    changed: Condvar,
18}
19
20static PROCESS_LOCAL_LOCKS: OnceLock<ProcessLocalLocks> = OnceLock::new();
21static NEXT_LOCAL_LOCK_OWNER: AtomicU64 = AtomicU64::new(1);
22
23pub(super) fn next_local_lock_owner() -> u64 {
24    NEXT_LOCAL_LOCK_OWNER.fetch_add(1, Ordering::Relaxed)
25}
26
27/// Provider-neutral distributed lock boundary.
28///
29/// Implementations must associate an acquired lock with `owner_token` and
30/// release it only while that token still owns the key. A zero timeout is one
31/// non-blocking attempt; a zero expiry means no automatic lease expiry.
32#[async_trait::async_trait]
33pub trait RemoteLockProvider: Send + Sync + 'static {
34    async fn try_remote_lock(
35        &self,
36        key: &str,
37        owner_token: &str,
38        timeout_millis: u64,
39        expire_millis: u64,
40    ) -> bool;
41
42    async fn unlock_remote(&self, key: &str, owner_token: &str) -> bool;
43}
44
45impl UserContext {
46    pub fn try_local_lock(&self, key: &str, timeout_millis: u64, expire_millis: u64) -> bool {
47        let locks = PROCESS_LOCAL_LOCKS.get_or_init(ProcessLocalLocks::default);
48        let deadline = Instant::now() + Duration::from_millis(timeout_millis);
49        let mut entries = locks.entries.lock().expect("local lock state poisoned");
50        loop {
51            let now = Instant::now();
52            match entries.get(key).copied() {
53                None => {
54                    entries.insert(
55                        key.to_owned(),
56                        LocalLockEntry {
57                            owner: self.local_lock_owner,
58                            expires_at: (expire_millis > 0)
59                                .then(|| now + Duration::from_millis(expire_millis)),
60                        },
61                    );
62                    return true;
63                }
64                Some(current)
65                    if current.owner == self.local_lock_owner
66                        || current.expires_at.is_some_and(|expiry| now >= expiry) =>
67                {
68                    entries.insert(
69                        key.to_owned(),
70                        LocalLockEntry {
71                            owner: self.local_lock_owner,
72                            expires_at: (expire_millis > 0)
73                                .then(|| now + Duration::from_millis(expire_millis)),
74                        },
75                    );
76                    return true;
77                }
78                Some(current) => {
79                    if timeout_millis == 0 || now >= deadline {
80                        return false;
81                    }
82                    let wake_after = current
83                        .expires_at
84                        .map(|expiry| expiry.saturating_duration_since(now))
85                        .unwrap_or_else(|| deadline.saturating_duration_since(now))
86                        .min(deadline.saturating_duration_since(now));
87                    let waited = locks
88                        .changed
89                        .wait_timeout(entries, wake_after)
90                        .expect("local lock state poisoned");
91                    entries = waited.0;
92                }
93            }
94        }
95    }
96
97    pub fn unlock_local(&self, key: &str) {
98        let locks = PROCESS_LOCAL_LOCKS.get_or_init(ProcessLocalLocks::default);
99        let mut entries = locks.entries.lock().expect("local lock state poisoned");
100        if entries
101            .get(key)
102            .is_some_and(|entry| entry.owner == self.local_lock_owner)
103        {
104            entries.remove(key);
105            locks.changed.notify_all();
106        }
107    }
108
109    /// Attempts to acquire a provider-backed distributed lock.
110    ///
111    /// A missing provider remains a no-op success, matching the optional
112    /// Remote Lock boundary in the other TeaQL runtimes. Install an
113    /// `Arc<dyn RemoteLockProvider>` resource to enable distributed exclusion.
114    pub async fn try_remote_lock(
115        &self,
116        key: &str,
117        timeout_millis: u64,
118        expire_millis: u64,
119    ) -> bool {
120        match self.get_resource::<Arc<dyn RemoteLockProvider>>() {
121            Some(provider) => {
122                provider
123                    .try_remote_lock(key, &self.remote_lock_owner, timeout_millis, expire_millis)
124                    .await
125            }
126            None => true,
127        }
128    }
129
130    /// Releases a distributed lock only when this context still owns it.
131    pub async fn unlock_remote(&self, key: &str) -> bool {
132        match self.get_resource::<Arc<dyn RemoteLockProvider>>() {
133            Some(provider) => provider.unlock_remote(key, &self.remote_lock_owner).await,
134            None => true,
135        }
136    }
137}
teaql_runtime/context/locking.rs

teaql_runtime/context/
locking.rs
