distkit 0.6.0

A toolkit of distributed systems primitives for Rust, backed by Redis
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
//! Distributed mutual-exclusion lock: [`Mutex`] and its RAII release guard
//! [`MutexGuard`].
//!
//! Mirrors the surface of [`tokio::sync::Mutex`] over Redis. The guard guards no
//! inner data — it is a pure release token. Acquire is fallible and async over
//! the network; release is best-effort on `Drop` plus an explicit awaitable
//! [`MutexGuard::release`]. A held lock renews its lease in the background every
//! `ttl/3`; a failed renewal marks the lease [`LockGuardState::Lost`], but the task
//! keeps retrying and clears the mark if ownership is later regained. The current
//! state is observable via [`MutexGuard::get_state`] and is also returned by
//! [`MutexGuard::release`].

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};

use redis::aio::ConnectionManager;
use tokio::task::JoinHandle;
use tokio::time::{MissedTickBehavior, interval};

use crate::DistkitError;
use crate::lock::{LockError, LockGuardState, LockOptions, mutex_backend};

/// A distributed mutual-exclusion lock backed by Redis.
///
/// One `Mutex` stands for exactly one resource: its Redis key and owner id are
/// fixed when you build it from [`LockOptions`], much like
/// `tokio::sync::Mutex::new(x)` binds the value it protects. The guard it hands
/// back protects no data of its own — it is purely a token that says "you hold
/// this lock right now."
///
/// # How a lock is held
///
/// Acquiring writes `key -> owner` in Redis with a TTL (the lease) and only
/// succeeds when nobody else holds the key, so you never receive a [`MutexGuard`]
/// without a confirmed acquisition. While the guard is alive a background task
/// renews the lease every `ttl / 3` so it never quietly expires under you.
/// Releasing — explicitly through [`MutexGuard::release`] or implicitly on
/// `Drop` — stops that task and deletes the key, but only while we still own it.
///
/// # When a held lock can still be lost
///
/// The lease lives in Redis, not in this process, so holding a guard is a strong
/// signal that you own the lock rather than an ironclad guarantee. Each refresh
/// asks Redis to confirm ownership and extend the lease. When a refresh can't be
/// confirmed — Redis is unreachable, the round-trip errors, or the key no longer
/// points at us — the lock is marked [`Lost`](LockGuardState::Lost). The refresh
/// task keeps going after that, and the next refresh it *can* confirm clears the
/// mark and returns the lock to [`Acquired`](LockGuardState::Acquired); a short
/// network blip usually heals itself this way.
///
/// A long partition is the case to watch. If we can't reach Redis for longer than
/// the TTL, the lease expires there and another owner is free to take the key —
/// at that point the lock really is gone, a later refresh has nothing to reclaim,
/// and it stays `Lost`. So if correctness depends on the lock, don't lean on the
/// guard's existence alone: call [`MutexGuard::get_state`] to re-check what we
/// currently believe before entering the critical section.
#[derive(Debug)]
pub struct Mutex {
    connection_manager: ConnectionManager,
    full_key: String,
    owner: String,
    ttl_ms: i64,
    ttl_duration: Duration,
    max_wait: Option<Duration>,
    retry_interval: Duration,
}

impl Mutex {
    /// Creates a new distributed mutex from the given options.
    ///
    /// The effective Redis key is `{namespace}:{key}`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use distkit::{DistkitRedisKey, lock::{Mutex, LockOptions}};
    ///
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let redis_url = std::env::var("REDIS_URL")
    ///     .unwrap_or_else(|_| "redis://127.0.0.1:6379".to_string());
    /// let client = redis::Client::open(redis_url)?;
    /// let conn = client.get_connection_manager().await?;
    /// let key = DistkitRedisKey::try_from("my_resource".to_string())?;
    /// let mutex = Mutex::new(LockOptions::new(key, conn));
    /// let guard = mutex.try_lock().await?;
    /// guard.release().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn new(options: LockOptions) -> Arc<Self> {
        let LockOptions {
            key,
            connection_manager,
            namespace,
            ttl,
            owner_id,
            max_wait,
            retry_interval,
            ..
        } = options;

        let full_key = format!("{}:{}", *namespace, *key);
        let owner = owner_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());

        Arc::new(Self {
            connection_manager,
            full_key,
            owner,
            ttl_ms: ttl.as_millis() as i64,
            ttl_duration: ttl,
            max_wait,
            retry_interval,
        })
    }

    /// Acquires the lock, waiting up to `max_wait` (or forever if `max_wait` is
    /// `None`), polling every `retry_interval`.
    pub async fn lock(&self) -> Result<MutexGuard, DistkitError> {
        self.acquire_core(self.max_wait, None, self.retry_interval)
            .await
    }

    /// Tries to acquire the lock in a single attempt without waiting. Returns
    /// [`LockError::AcquireFail`] if the lock is already held.
    pub async fn try_lock(&self) -> Result<MutexGuard, DistkitError> {
        self.acquire_core(Some(Duration::ZERO), None, Duration::ZERO)
            .await
    }

    /// Tries to acquire the lock, waiting up to `timeout` and polling at the
    /// lock's configured `retry_interval` (see [`LockOptions::retry_interval`]).
    /// Returns [`LockError::Timeout`] if the deadline passes first.
    pub async fn try_lock_with_timeout(&self, timeout: Duration) -> Result<MutexGuard, DistkitError> {
        self.acquire_core(Some(timeout), None, self.retry_interval)
            .await
    }

    /// Tries to acquire the lock, making up to `max_retries` retries after the
    /// initial attempt (`max_retries + 1` attempts total), polling at the lock's
    /// configured `retry_interval` (see [`LockOptions::retry_interval`]). Returns
    /// [`LockError::RetriesExhausted`] if every attempt fails.
    ///
    /// Relies on a non-zero configured `retry_interval` (the default); a lock
    /// configured with a zero `retry_interval` collapses to a single attempt.
    pub async fn try_lock_with_retries(
        &self,
        max_retries: usize,
    ) -> Result<MutexGuard, DistkitError> {
        self.acquire_core(None, Some(max_retries), self.retry_interval)
            .await
    }

    /// Tries to acquire the lock, waiting up to `timeout` and polling every
    /// `retry_interval`. Returns [`LockError::Timeout`] if the deadline passes
    /// first. A `retry_interval` of zero is a tight spin.
    #[deprecated(
        since = "0.5.3",
        note = "use `try_lock_with_timeout`, which sources the retry interval from the lock's configuration"
    )]
    pub async fn try_lock_for(
        &self,
        timeout: Duration,
        retry_interval: Duration,
    ) -> Result<MutexGuard, DistkitError> {
        self.acquire_core(Some(timeout), None, retry_interval).await
    }

    /// Shared acquire retry-loop backing every public form.
    ///
    /// `timeout == None` waits forever; `Some(ZERO)` is a single shot;
    /// `Some(duration)` is a bounded wait. `max_retries == Some(n)` gives up with
    /// [`LockError::RetriesExhausted`] after `n` retries; `None` is unbounded by
    /// retry count.
    async fn acquire_core(
        &self,
        timeout: Option<Duration>,
        max_retries: Option<usize>,
        retry_interval: Duration,
    ) -> Result<MutexGuard, DistkitError> {
        let start = Instant::now();
        let mut connection = self.connection_manager.clone();
        let mut retry_interval = if retry_interval.is_zero() {
            None
        } else {
            let mut retry_interval = interval(retry_interval);
            retry_interval.set_missed_tick_behavior(MissedTickBehavior::Delay);

            Some(retry_interval)
        };

        let mut attempt = 0usize;
        loop {
            if let Some(retry_interval) = &mut retry_interval {
                // First tick is immediately, subsequent ticks are delayed.
                retry_interval.tick().await;
            }

            let acquired =
                mutex_backend::acquire(&mut connection, &self.full_key, &self.owner, self.ttl_ms)
                    .await?;

            if acquired {
                let lost = Arc::new(AtomicBool::new(false));
                let refresh_handle = self.spawn_refresh(lost.clone());

                return Ok(MutexGuard {
                    connection_manager: self.connection_manager.clone(),
                    full_key: self.full_key.clone(),
                    owner: self.owner.clone(),
                    refresh_handle: Some(refresh_handle),
                    lost,
                    on_attempt: attempt,
                });
            }

            if retry_interval.is_none() {
                return Err(LockError::AcquireFail.into());
            }

            if let Some(ttl) = timeout {
                if ttl.is_zero() {
                    return Err(LockError::AcquireFail.into());
                }

                let waited = start.elapsed();
                if waited >= ttl {
                    return Err(LockError::Timeout { waited }.into());
                }
            }

            if let Some(max_retries) = max_retries
                && attempt >= max_retries
            {
                return Err(LockError::RetriesExhausted { retries: attempt }.into());
            }

            attempt += 1;
        }
    }

    /// Spawns the background lease-renewal task for a held lock.
    ///
    /// Ticks every `ttl/3`, refreshing the lease while we still own it. The first
    /// (immediate) tick is skipped since `acquire` just set the lease. On any failed
    /// renewal — lost ownership (`Ok(false)`) or a transport error (`Err`) — it sets
    /// the shared `lost` flag but keeps ticking; if a later refresh succeeds it
    /// clears the flag again. The flag is surfaced via [`MutexGuard::get_state`] /
    /// [`MutexGuard::release`] as [`LockGuardState::Lost`]. The task runs until the guard
    /// aborts it on release or drop.
    fn spawn_refresh(&self, lost: Arc<AtomicBool>) -> JoinHandle<()> {
        let mut connection_manager = self.connection_manager.clone();
        let full_key = self.full_key.clone();
        let owner = self.owner.clone();
        let ttl_ms = self.ttl_ms;
        let ttl_duration = self.ttl_duration;

        tokio::spawn(async move {
            let mut ticker = interval(ttl_duration / 3);
            ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);

            // Skip the first immediate tick — acquire just set the lease.
            ticker.tick().await;

            loop {
                ticker.tick().await;

                match mutex_backend::refresh(&mut connection_manager, &full_key, &owner, ttl_ms)
                    .await
                {
                    Ok(true) => {
                        // Refresh succeeded; if the lease had been marked lost, we
                        // just regained ownership — clear the flag.
                        if lost.swap(false, Ordering::AcqRel) {
                            tracing::debug!(
                                full_key,
                                owner,
                                "Lost distributed lock reaquired during refresh"
                            );
                        }
                    }
                    Ok(false) => {
                        tracing::debug!(
                            full_key,
                            owner,
                            "Lost distributed lock lease during refresh"
                        );
                        lost.store(true, Ordering::Release);
                    }
                    Err(error) => {
                        tracing::debug!(
                            ?error,
                            full_key,
                            owner,
                            "Error refreshing distributed lock lease"
                        );
                        lost.store(true, Ordering::Release);
                    }
                }
            }
        })
    } // end spawn_refresh
}

/// A token proving the current task holds a [`Mutex`].
///
/// You only ever get a `MutexGuard` from a confirmed acquisition, so its very
/// existence means the lock was yours at acquire time. It stays good for as long
/// as the background refresh keeps renewing the lease — see [`Mutex`] for the
/// cases where that can fail and the lock turns up [`Lost`](LockGuardState::Lost).
///
/// Dropping the guard releases the lock on a best-effort, fire-and-forget basis.
/// Reach for [`release`](MutexGuard::release) when you want to wait for the
/// release to land and learn the final [`LockGuardState`], and for
/// [`get_state`](MutexGuard::get_state) when you want to re-check ownership while
/// still holding it. The guard carries no inner data.
#[derive(Debug)]
pub struct MutexGuard {
    connection_manager: ConnectionManager,
    full_key: String,
    owner: String,
    refresh_handle: Option<JoinHandle<()>>,
    lost: Arc<AtomicBool>,
    /// Zero-based index of the acquire attempt that obtained this lock.
    on_attempt: usize,
}

impl MutexGuard {
    /// Re-checks what we currently believe the lock's state to be.
    ///
    /// Returns [`Acquired`](LockGuardState::Acquired) while the lease is being
    /// renewed normally, [`Lost`](LockGuardState::Lost) when a recent refresh
    /// couldn't confirm ownership (see [`Mutex`] for why that happens and when it
    /// recovers), and [`Released`](LockGuardState::Released) once the guard has
    /// been released.
    ///
    /// This reads the latest result the background refresh recorded — it does not
    /// make its own Redis round-trip. Reach for it before a critical section when
    /// you need more confidence than "I'm still holding the guard."
    pub fn get_state(&self) -> LockGuardState {
        if self.refresh_handle.is_none() {
            return LockGuardState::Released;
        }

        if self.lost.load(Ordering::Acquire) {
            return LockGuardState::Lost;
        }

        LockGuardState::Acquired
    }

    /// The zero-based acquire attempt that obtained this lock: `0` if the very
    /// first poll succeeded, `n` after `n` retries. A one-shot `try_lock` is
    /// always `0`. Useful for contention metrics.
    pub fn get_on_attempt(&self) -> usize {
        self.on_attempt
    }

    /// Releases the lock and reports its final [`LockGuardState`].
    ///
    /// Stops the background refresh, then deletes the key if we still own it,
    /// awaiting the round-trip so the caller can act on the outcome. Returns
    /// [`Released`](LockGuardState::Released) on a clean release, or
    /// [`Lost`](LockGuardState::Lost) if the lease had already slipped away — in
    /// that case no delete is issued, since the key may now belong to another
    /// owner. A failed Redis round-trip surfaces as `Err`.
    pub async fn release(mut self) -> Result<LockGuardState, DistkitError> {
        if let Some(handle) = self.refresh_handle.take() {
            handle.abort();
        }

        if self.lost.load(Ordering::Acquire) {
            return Ok(LockGuardState::Lost);
        }

        let mut connection = self.connection_manager.clone();
        mutex_backend::release(&mut connection, &self.full_key, &self.owner).await?;

        Ok(LockGuardState::Released)
    }
}

impl Drop for MutexGuard {
    fn drop(&mut self) {
        let Some(handle) = self.refresh_handle.take() else {
            // Already released
            return;
        };

        handle.abort();

        let mut connection = self.connection_manager.clone();
        let full_key = self.full_key.clone();
        let owner = self.owner.clone();

        tokio::spawn(async move {
            if let Err(error) = mutex_backend::release(&mut connection, &full_key, &owner).await {
                tracing::error!(?error, full_key, "Error releasing distributed lock on drop");
            }
        });
    }
}