cachekit/flight.rs
1//! Cold-miss single-flight: dedup concurrent fills of the same key.
2//!
3//! Under metered-misses pricing a stampede is literally billable — N tasks
4//! missing the same key at once means N backend misses and N executions of
5//! the wrapped function. [`CacheKit::single_flight`](crate::CacheKit::single_flight)
6//! collapses that to one:
7//!
8//! - **In-process** (always available): a per-key async mutex. The first
9//! task through becomes the *leader* and computes; concurrent tasks queue
10//! behind it and re-check the cache once the leader finishes.
11//! - **Cross-process** (`reliability` feature, native, backend implements
12//! `LockableBackend` — CachekitIO and Redis do): the leader additionally
13//! takes a distributed fill lock. If another process already holds it,
14//! this process polls the cache for the other side's fill instead of
15//! recomputing, and computes anyway once the lock timeout, measured from
16//! its own lock attempt, has passed: by then the other side's lock has
17//! expired (fail-open — a stampede beats unavailability).
18//!
19//! Once a leader's fill is stored, the `#[cachekit]` macro sends its unlock
20//! from a task spawned on the ambient tokio runtime and returns without
21//! waiting for it. That needs `l1` and a tokio runtime, and never happens
22//! under `unsync`; there, and whenever the fill was not stored, the caller
23//! awaits the unlock. If the backend has not answered the spawned unlock by
24//! the next lock attempt on that key from this client or its clones (still
25//! in flight, or its runtime idle or dropped), that attempt sends the unlock
26//! itself first, so a re-miss does not contest its own lock unless that
27//! unlock fails too. If the runtime is dropped or left idle right after the
28//! call, the unlock waits for that next attempt, and until then the lease
29//! stays held server-side for up to its 5 s timeout: another process that
30//! misses the key inside that window polls for it, as it would behind any
31//! slow leader.
32//!
33//! The `#[cachekit]` macro wires this in automatically around its miss path.
34//! Its stale-while-revalidate refresh takes the same locks but never waits
35//! for another worker's fill: if another worker holds either one, or the lock
36//! call fails, the refresh stands down without polling or re-reading the
37//! cache, and the stale copy keeps being served. A cold miss queued behind a
38//! holder that found the distributed lock contested and let go without a
39//! fill (a refresh that stood down, a cancelled or failed poll) contests that
40//! lock itself before computing, since another process may be filling.
41//! Manual usage follows the same shape:
42//!
43//! ```no_run
44//! # async fn example(cache: &cachekit::CacheKit) -> Result<(), cachekit::CachekitError> {
45//! if let Some(_v) = cache.get::<String>("expensive").await? {
46//! return Ok(());
47//! }
48//! let mut flight = cache.single_flight("expensive").await;
49//! while flight.wait_for_fill().await {
50//! if let Some(_v) = cache.get::<String>("expensive").await? {
51//! flight.release().await; // another worker filled it
52//! return Ok(());
53//! }
54//! }
55//! let value = "computed".to_owned(); // expensive work — runs once
56//! cache.set("expensive", &value).await?;
57//! flight.release().await;
58//! # Ok(())
59//! # }
60//! ```
61
62use std::collections::HashMap;
63use std::sync::{Arc, Mutex, PoisonError, Weak};
64
65#[cfg(all(
66 feature = "reliability",
67 feature = "l1",
68 not(feature = "unsync"),
69 not(target_arch = "wasm32")
70))]
71use std::sync::atomic::AtomicBool;
72#[cfg(all(feature = "l1", not(feature = "unsync"), not(target_arch = "wasm32")))]
73use std::sync::atomic::{AtomicU64, Ordering};
74
75use crate::client::SharedBackend;
76
77/// Above this many live entries, dead map slots are swept opportunistically.
78const SWEEP_THRESHOLD: usize = 128;
79
80/// How long a distributed fill lock is held server-side before auto-expiry.
81///
82/// Also the cross-process suppression ceiling: a fill that runs longer than
83/// this loses its lock mid-compute and another process may recompute
84/// concurrently (fail-open by design — release is owner-checked, so an
85/// expired lock is never wrongfully deleted). Workloads whose fills
86/// routinely approach 5 s keep in-process dedup but should not rely on
87/// cross-process suppression.
88/// ponytail: fixed TTL, no heartbeat — add lock renewal if slow fills matter.
89#[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
90const FILL_LOCK_TIMEOUT_MS: u64 = 5_000;
91
92/// Poll cadence while waiting for another process's fill.
93#[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
94const FILL_POLL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(100);
95
96// ── FlightMap ────────────────────────────────────────────────────────────────
97
98/// Per-key async mutexes for in-process fill dedup. Weak entries let finished
99/// flights drop their state without an explicit removal protocol.
100///
101/// Each mutex guards a stand-down note: `true` when the last holder to try
102/// the distributed lock found it contested, so it may have filled nothing (a
103/// refresh stands down; a cold miss polls and may be cancelled or fail), and
104/// a queued cold miss whose re-check misses must contest that lock before
105/// computing. A holder that leads clears the note: on acquire, or when its
106/// own lease contest makes it a leader. A follower leaves it as it found it,
107/// whether it finds a fill, errors out or is dropped, so the next follower
108/// still sees it. A stale `true` costs at most one lock call, by a follower
109/// that was about to compute anyway.
110#[derive(Default)]
111pub(crate) struct FlightMap {
112 entries: Mutex<Entries>,
113}
114
115#[derive(Default)]
116struct Entries {
117 slots: HashMap<String, Slot>,
118 /// Sweep dead slots once the map grows past this: twice what the last
119 /// sweep left, so a large live set does not make every miss O(n).
120 sweep_above: usize,
121 /// When the last sweep ran. A sweep during a cold-key burst counts the
122 /// burst's flights as survivors, so the doubling alone would keep their
123 /// dead slots until the map doubles again; a sweep older than one lease
124 /// timeout, when every pin it saw has expired, is redone on the next miss.
125 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
126 swept_at: Option<tokio::time::Instant>,
127}
128
129impl Entries {
130 fn sweep_due(&self) -> bool {
131 let len = self.slots.len();
132 if len <= SWEEP_THRESHOLD {
133 return false;
134 }
135 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
136 if self
137 .swept_at
138 .is_none_or(|at| at.elapsed() >= std::time::Duration::from_millis(FILL_LOCK_TIMEOUT_MS))
139 {
140 return true;
141 }
142 len > self.sweep_above
143 }
144
145 fn sweep(&mut self) {
146 self.slots.retain(|_, slot| {
147 #[cfg(all(
148 feature = "reliability",
149 feature = "l1",
150 not(feature = "unsync"),
151 not(target_arch = "wasm32")
152 ))]
153 slot.expire_pin();
154 slot.entry.strong_count() > 0
155 });
156 self.sweep_above = 2 * self.slots.len();
157 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
158 {
159 self.swept_at = Some(tokio::time::Instant::now());
160 }
161 }
162}
163
164struct Slot {
165 entry: Weak<tokio::sync::Mutex<KeyState>>,
166 /// Keeps the entry, and with it a pending unlock, alive after a leader
167 /// detached its unlock, so a re-miss finds the unlock even when its task
168 /// was cancelled with its runtime: until the backend answers the unlock
169 /// or, failing that, at least one lease timeout. An expired pin goes at
170 /// the next sweep.
171 #[cfg(all(
172 feature = "reliability",
173 feature = "l1",
174 not(feature = "unsync"),
175 not(target_arch = "wasm32")
176 ))]
177 pin: Option<Pin>,
178}
179
180#[cfg(all(
181 feature = "reliability",
182 feature = "l1",
183 not(feature = "unsync"),
184 not(target_arch = "wasm32")
185))]
186struct Pin {
187 until: tokio::time::Instant,
188 /// The pending unlock this pin is for (see [`PendingRelease::answered`]).
189 answered: Arc<AtomicBool>,
190 _entry: Arc<tokio::sync::Mutex<KeyState>>,
191}
192
193#[cfg(all(
194 feature = "reliability",
195 feature = "l1",
196 not(feature = "unsync"),
197 not(target_arch = "wasm32")
198))]
199impl Slot {
200 fn expire_pin(&mut self) {
201 if self
202 .pin
203 .as_ref()
204 .is_some_and(|pin| pin.until <= tokio::time::Instant::now())
205 {
206 self.pin = None;
207 }
208 }
209}
210
211/// What the per-key mutex in [`FlightMap`] guards.
212#[derive(Default)]
213struct KeyState {
214 /// The stand-down note (see [`FlightMap`]).
215 stand_down: bool,
216 /// An unlock spawned by a leader that stored its fill and returned
217 /// without waiting for it ([`SingleFlight::__release_detached`]).
218 // `l1` stands in for tokio's `rt` feature (spawn), which `reliability`
219 // alone does not enable.
220 #[cfg(all(
221 feature = "reliability",
222 feature = "l1",
223 not(feature = "unsync"),
224 not(target_arch = "wasm32")
225 ))]
226 pending_release: Option<PendingRelease>,
227}
228
229/// A detached unlock. Holds no handle to its task: that would keep the
230/// task's runtime (and its file descriptors) alive after the runtime is
231/// dropped.
232#[cfg(all(
233 feature = "reliability",
234 feature = "l1",
235 not(feature = "unsync"),
236 not(target_arch = "wasm32")
237))]
238struct PendingRelease {
239 lock_id: String,
240 /// Set once the backend answered the unlock (released, or already gone).
241 answered: Arc<AtomicBool>,
242}
243
244#[cfg(all(
245 feature = "reliability",
246 feature = "l1",
247 not(feature = "unsync"),
248 not(target_arch = "wasm32")
249))]
250impl PendingRelease {
251 /// Make sure the unlock has reached the backend before this client
252 /// contests the key again: a re-miss would otherwise find its own lease
253 /// held and poll an empty cache until the lock timeout. Never waits on
254 /// the spawned task, whose runtime may be idle or gone: an unanswered
255 /// unlock is sent again from here. `LockableBackend::release_lock` is
256 /// owner-checked, so a duplicate, or the task's late original, is
257 /// harmless. If this unlock fails too, the re-miss polls like any
258 /// contested follower, for at most one lock timeout.
259 async fn settle(&self, backend: &SharedBackend, full_key: &str) {
260 if !self.answered.load(Ordering::Acquire) {
261 release_lock(backend, full_key, &self.lock_id).await;
262 }
263 }
264}
265
266impl FlightMap {
267 fn handle(&self, key: &str) -> Arc<tokio::sync::Mutex<KeyState>> {
268 let mut entries = self.entries.lock().unwrap_or_else(PoisonError::into_inner);
269 // ponytail: O(n) sweep, amortised by doubling and at most one per
270 // lease timeout otherwise; a doubly-indexed structure is not worth it
271 // until someone caches millions of distinct cold keys.
272 if entries.sweep_due() {
273 entries.sweep();
274 }
275 if let Some(existing) = entries.slots.get(key).and_then(|slot| slot.entry.upgrade()) {
276 return existing;
277 }
278 let fresh = Arc::new(tokio::sync::Mutex::new(KeyState::default()));
279 entries.slots.insert(
280 key.to_owned(),
281 Slot {
282 entry: Arc::downgrade(&fresh),
283 #[cfg(all(
284 feature = "reliability",
285 feature = "l1",
286 not(feature = "unsync"),
287 not(target_arch = "wasm32")
288 ))]
289 pin: None,
290 },
291 );
292 fresh
293 }
294
295 /// Pin `entry` (the entry for `key`) for one lease timeout, on behalf of
296 /// the pending unlock flagged by `answered`.
297 #[cfg(all(
298 feature = "reliability",
299 feature = "l1",
300 not(feature = "unsync"),
301 not(target_arch = "wasm32")
302 ))]
303 fn pin(&self, key: &str, entry: Arc<tokio::sync::Mutex<KeyState>>, answered: Arc<AtomicBool>) {
304 let until =
305 tokio::time::Instant::now() + std::time::Duration::from_millis(FILL_LOCK_TIMEOUT_MS);
306 let mut entries = self.entries.lock().unwrap_or_else(PoisonError::into_inner);
307 if let Some(slot) = entries.slots.get_mut(key) {
308 slot.pin = Some(Pin {
309 until,
310 answered,
311 _entry: entry,
312 });
313 }
314 }
315
316 /// Drop `key`'s pin if it is still the one for the unlock flagged by
317 /// `answered`: a later leader's pin for the same key stays.
318 #[cfg(all(
319 feature = "reliability",
320 feature = "l1",
321 not(feature = "unsync"),
322 not(target_arch = "wasm32")
323 ))]
324 fn unpin(&self, key: &str, answered: &Arc<AtomicBool>) {
325 let mut entries = self.entries.lock().unwrap_or_else(PoisonError::into_inner);
326 if let Some(slot) = entries.slots.get_mut(key) {
327 if slot
328 .pin
329 .as_ref()
330 .is_some_and(|pin| Arc::ptr_eq(&pin.answered, answered))
331 {
332 slot.pin = None;
333 }
334 }
335 }
336}
337
338// ── SWR mutation ordering ────────────────────────────────────────────────────
339
340/// State shared by a stale-read token and same-key mutations while that token
341/// is alive. The weak map can discard idle keys without losing correctness:
342/// an outstanding token itself keeps this state alive.
343#[cfg(all(feature = "l1", not(feature = "unsync"), not(target_arch = "wasm32")))]
344pub(crate) struct MutationState {
345 lock: Arc<tokio::sync::Mutex<()>>,
346 version: AtomicU64,
347}
348
349#[cfg(all(feature = "l1", not(feature = "unsync"), not(target_arch = "wasm32")))]
350impl MutationState {
351 fn new() -> Self {
352 Self {
353 lock: Arc::new(tokio::sync::Mutex::new(())),
354 version: AtomicU64::new(0),
355 }
356 }
357
358 pub(crate) fn version(&self) -> u64 {
359 self.version.load(Ordering::Acquire)
360 }
361}
362
363/// Per-key mutation versions for conditional SWR commits. Cloned clients
364/// share this map. Each entry is weak so unrelated keys do not accumulate.
365#[cfg(all(feature = "l1", not(feature = "unsync"), not(target_arch = "wasm32")))]
366#[derive(Default)]
367pub(crate) struct MutationMap {
368 entries: Mutex<HashMap<String, Weak<MutationState>>>,
369}
370
371#[cfg(all(feature = "l1", not(feature = "unsync"), not(target_arch = "wasm32")))]
372impl MutationMap {
373 pub(crate) fn state(&self, key: &str) -> Arc<MutationState> {
374 let mut map = self.entries.lock().unwrap_or_else(PoisonError::into_inner);
375 if map.len() > SWEEP_THRESHOLD {
376 map.retain(|_, weak| weak.strong_count() > 0);
377 }
378 if let Some(existing) = map.get(key).and_then(Weak::upgrade) {
379 return existing;
380 }
381 let fresh = Arc::new(MutationState::new());
382 map.insert(key.to_owned(), Arc::downgrade(&fresh));
383 fresh
384 }
385
386 pub(crate) async fn lock(&self, key: &str) -> MutationGuard {
387 let state = self.state(key);
388 let lock = Arc::clone(&state.lock).lock_owned().await;
389 MutationGuard { state, _lock: lock }
390 }
391}
392
393/// Holds same-key ordering across an L2 mutation and its L1 update.
394#[cfg(all(feature = "l1", not(feature = "unsync"), not(target_arch = "wasm32")))]
395pub(crate) struct MutationGuard {
396 state: Arc<MutationState>,
397 _lock: tokio::sync::OwnedMutexGuard<()>,
398}
399
400#[cfg(all(feature = "l1", not(feature = "unsync"), not(target_arch = "wasm32")))]
401impl MutationGuard {
402 pub(crate) fn snapshot(&self) -> (Arc<MutationState>, u64) {
403 (Arc::clone(&self.state), self.state.version())
404 }
405
406 pub(crate) fn is_current(&self, state: &Arc<MutationState>, version: u64) -> bool {
407 Arc::ptr_eq(&self.state, state) && self.state.version() == version
408 }
409
410 pub(crate) fn advance(&self) {
411 self.state.version.fetch_add(1, Ordering::Release);
412 }
413}
414
415// ── SingleFlight guard ───────────────────────────────────────────────────────
416
417enum Role {
418 /// First worker in: compute without re-checking (a re-check would be a
419 /// second billable miss under metered-misses pricing).
420 Leader,
421 /// Queued behind a local holder that has since finished: re-check the
422 /// cache once — a leader's fill is in L1 — then compute if it missed.
423 /// `contest_lease` is the stand-down note this follower found (see
424 /// [`FlightMap`]): another process may be filling, so a missed re-check
425 /// contests the lease before computing.
426 LocalFollower {
427 rechecked: bool,
428 #[cfg_attr(
429 not(all(feature = "reliability", not(target_arch = "wasm32"))),
430 allow(dead_code)
431 )]
432 contest_lease: bool,
433 },
434 /// Another *process* holds the distributed fill lock: poll the cache for
435 /// its fill, then compute anyway at `deadline` (fail-open). The deadline
436 /// is the lock timeout after this flight's own lock attempt, so the other
437 /// process's lock, taken earlier, has expired by then.
438 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
439 RemoteContested { deadline: tokio::time::Instant },
440}
441
442/// Outcome of one distributed fill-lock attempt. The cold-miss and refresh
443/// paths apply different policies to it.
444#[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
445enum LockAttempt {
446 /// The lease was granted; carries its id.
447 Granted(String),
448 /// The backend has no distributed lock.
449 Unlockable,
450 /// Another process holds the lease, which has expired by `deadline`:
451 /// one lock timeout after this attempt started.
452 Contested { deadline: tokio::time::Instant },
453 /// The lock call itself failed.
454 Failed,
455}
456
457#[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
458impl LockAttempt {
459 /// Must be called with the key's in-process lock held (`state`).
460 async fn run(state: &mut KeyState, backend: &SharedBackend, full_key: &str) -> Self {
461 let Some(lockable) = backend.as_lockable() else {
462 return Self::Unlockable;
463 };
464 // Cleared only once settled, so a cancelled settle is redone by the
465 // next attempt.
466 #[cfg(all(feature = "l1", not(feature = "unsync")))]
467 if let Some(pending) = &state.pending_release {
468 pending.settle(backend, full_key).await;
469 state.pending_release = None;
470 }
471 #[cfg(not(all(feature = "l1", not(feature = "unsync"))))]
472 let _ = state;
473 let timeout = std::time::Duration::from_millis(FILL_LOCK_TIMEOUT_MS);
474 let deadline = tokio::time::Instant::now() + timeout;
475 match lockable.acquire_lock(full_key, FILL_LOCK_TIMEOUT_MS).await {
476 Ok(Some(lock_id)) => Self::Granted(lock_id),
477 Ok(None) => Self::Contested { deadline },
478 Err(err) => {
479 crate::metrics::trace_lock_error(full_key, "lock", &err);
480 Self::Failed
481 }
482 }
483 }
484
485 /// Cold-miss policy: poll a contested lease; a lock-infrastructure error
486 /// fails open to a plain leader — suppression is an optimisation, never
487 /// an availability dependency.
488 fn cold_role(self) -> (Role, Option<String>) {
489 match self {
490 Self::Granted(lock_id) => (Role::Leader, Some(lock_id)),
491 Self::Unlockable | Self::Failed => (Role::Leader, None),
492 Self::Contested { deadline } => (Role::RemoteContested { deadline }, None),
493 }
494 }
495}
496
497/// Guard for a single-flight fill, returned by
498/// [`CacheKit::single_flight`](crate::CacheKit::single_flight).
499///
500/// Holds the per-key in-process lock for its whole lifetime, and the
501/// distributed fill lock (if one was acquired) until [`Self::release`].
502/// Dropping without `release` is safe: the in-process lock frees immediately
503/// and a distributed lock expires server-side after its timeout.
504pub struct SingleFlight {
505 /// The per-key lock, guarding the key's [`KeyState`].
506 local: tokio::sync::OwnedMutexGuard<KeyState>,
507 /// The map `local` came from, to pin it while an unlock is pending.
508 #[cfg(all(
509 feature = "reliability",
510 feature = "l1",
511 not(feature = "unsync"),
512 not(target_arch = "wasm32")
513 ))]
514 map: crate::client::SharedFlight,
515 role: Role,
516 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
517 backend: SharedBackend,
518 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
519 full_key: String,
520 /// Id of the distributed fill lock this flight holds, if any.
521 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
522 lock_id: Option<String>,
523}
524
525impl SingleFlight {
526 /// `true` while another worker may still be filling this key — re-check
527 /// the cache after every `true` before computing yourself:
528 ///
529 /// - Leader: immediately `false` (compute, don't re-read your own miss).
530 /// - Queued behind a local holder: `true` once, for the re-check, then
531 /// `false`. If the holder found the distributed lock contested and let
532 /// go without a fill, a missed re-check first contests that lock and
533 /// then behaves as a leader or a contested flight.
534 /// - Contested cross-process: sleeps one poll interval per call, `true`
535 /// until the lock timeout (5 s), measured from this flight's lock
536 /// attempt, has passed. A hung re-check is not cut short here;
537 /// `#[cachekit]` bounds its own.
538 pub async fn wait_for_fill(&mut self) -> bool {
539 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
540 if let Role::LocalFollower {
541 rechecked: true,
542 contest_lease: true,
543 } = self.role
544 {
545 (self.role, self.lock_id) =
546 LockAttempt::run(&mut self.local, &self.backend, &self.full_key)
547 .await
548 .cold_role();
549 self.note_role();
550 }
551 match &mut self.role {
552 Role::Leader => false,
553 Role::LocalFollower { rechecked, .. } => {
554 let first = !*rechecked;
555 *rechecked = true;
556 first
557 }
558 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
559 Role::RemoteContested { deadline } => {
560 let deadline = *deadline;
561 let now = tokio::time::Instant::now();
562 if now >= deadline {
563 return false;
564 }
565 // The last sleep stops at the deadline, and no re-check
566 // follows it: the other side's lock has expired by then.
567 tokio::time::sleep(FILL_POLL_INTERVAL.min(deadline - now)).await;
568 tokio::time::Instant::now() < deadline
569 }
570 }
571 }
572
573 /// Release the flight. Best-effort: frees the distributed fill lock (if
574 /// held) so other processes stop waiting early. A failed release is
575 /// traced (feature `tracing`) and otherwise ignored: the lock expires
576 /// server-side regardless. Leaves the stand-down note as it is (see
577 /// [`FlightMap`]), since a caller may release after an error, with no fill.
578 pub async fn release(self) {
579 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
580 if let Some(lock_id) = &self.lock_id {
581 release_lock(&self.backend, &self.full_key, lock_id).await;
582 }
583 }
584
585 /// Release after the fill was stored, without waiting for the unlock
586 /// round trip. Macro plumbing for `#[cachekit]` — not public API.
587 ///
588 /// The unlock is spawned onto the ambient tokio runtime and the per-key
589 /// lock frees at once, so the caller and its same-key local followers
590 /// return one round trip sooner. Remote waiters poll the cache, not the
591 /// lock, and the fill is already stored. Only call it after a stored
592 /// fill: after an error or a failed store, a quick re-call would miss
593 /// again, so [`Self::release`] keeps the unlock on the caller's path.
594 /// Without `l1`, under `unsync` or off a tokio runtime this is
595 /// [`Self::release`]. A later lock attempt on the key re-sends an unlock
596 /// the backend has not answered (see [`PendingRelease::settle`]).
597 #[doc(hidden)]
598 #[cfg_attr(
599 not(all(
600 feature = "reliability",
601 feature = "l1",
602 not(feature = "unsync"),
603 not(target_arch = "wasm32")
604 )),
605 allow(unused_mut)
606 )]
607 pub async fn __release_detached(mut self) {
608 #[cfg(all(
609 feature = "reliability",
610 feature = "l1",
611 not(feature = "unsync"),
612 not(target_arch = "wasm32")
613 ))]
614 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
615 if let Some(lock_id) = self.lock_id.take() {
616 let answered = Arc::new(AtomicBool::new(false));
617 let full_key = std::mem::take(&mut self.full_key);
618 self.map.pin(
619 &full_key,
620 Arc::clone(tokio::sync::OwnedMutexGuard::mutex(&self.local)),
621 Arc::clone(&answered),
622 );
623 let unlock = {
624 let (map, backend, lock_id, answered) = (
625 self.map.clone(),
626 self.backend.clone(),
627 lock_id.clone(),
628 Arc::clone(&answered),
629 );
630 async move {
631 if release_lock(&backend, &full_key, &lock_id).await {
632 answered.store(true, Ordering::Release);
633 map.unpin(&full_key, &answered);
634 }
635 }
636 };
637 // Keep the caller's span, so a failed unlock traces under it.
638 #[cfg(feature = "tracing")]
639 let unlock = tracing::Instrument::instrument(unlock, tracing::Span::current());
640 drop(runtime.spawn(unlock));
641 self.local.pending_release = Some(PendingRelease { lock_id, answered });
642 }
643 return;
644 }
645 self.release().await;
646 }
647
648 /// Run one poll re-check, cut off at a contested flight's deadline.
649 /// `None` means the deadline passed first; treat it as a miss.
650 /// Macro plumbing for `#[cachekit]` — not public API.
651 #[doc(hidden)]
652 pub async fn __bounded_poll<F: std::future::Future>(&self, read: F) -> Option<F::Output> {
653 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
654 if let Role::RemoteContested { deadline } = self.role {
655 return tokio::time::timeout_at(deadline, read).await.ok();
656 }
657 Some(read.await)
658 }
659
660 /// Record this flight's role in the stand-down note (see [`FlightMap`]):
661 /// a leader clears it, a contested flight sets it, since it fills nothing
662 /// unless its poll runs out, and a follower leaves it as found.
663 fn note_role(&mut self) {
664 match self.role {
665 Role::Leader => self.local.stand_down = false,
666 Role::LocalFollower { .. } => {}
667 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
668 Role::RemoteContested { .. } => self.local.stand_down = true,
669 }
670 }
671
672 /// A flight over a freshly acquired per-key lock, its role recorded in
673 /// the stand-down note.
674 fn new(
675 map: &crate::client::SharedFlight,
676 local: tokio::sync::OwnedMutexGuard<KeyState>,
677 role: Role,
678 backend: &SharedBackend,
679 full_key: &str,
680 ) -> Self {
681 #[cfg(not(all(feature = "reliability", not(target_arch = "wasm32"))))]
682 let _ = (backend, full_key);
683 #[cfg(not(all(
684 feature = "reliability",
685 feature = "l1",
686 not(feature = "unsync"),
687 not(target_arch = "wasm32")
688 )))]
689 let _ = map;
690 let mut flight = Self {
691 local,
692 #[cfg(all(
693 feature = "reliability",
694 feature = "l1",
695 not(feature = "unsync"),
696 not(target_arch = "wasm32")
697 ))]
698 map: map.clone(),
699 role,
700 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
701 backend: backend.clone(),
702 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
703 full_key: full_key.to_owned(),
704 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
705 lock_id: None,
706 };
707 flight.note_role();
708 flight
709 }
710
711 /// Cold-miss entry: lead (attempting cross-process suppression via the
712 /// backend's distributed lock, when available), or queue behind a local
713 /// holder.
714 pub(crate) async fn acquire(
715 map: &crate::client::SharedFlight,
716 backend: &SharedBackend,
717 full_key: &str,
718 ) -> Self {
719 let handle = map.handle(full_key);
720 #[cfg_attr(
721 not(all(feature = "reliability", not(target_arch = "wasm32"))),
722 allow(unused_mut)
723 )]
724 let Ok(mut local) = Arc::clone(&handle).try_lock_owned() else {
725 // Contended: a local holder is filling. Queue behind it.
726 let local = handle.lock_owned().await;
727 let role = Role::LocalFollower {
728 rechecked: false,
729 contest_lease: local.stand_down,
730 };
731 return Self::new(map, local, role, backend, full_key);
732 };
733 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
734 {
735 let (role, lock_id) = LockAttempt::run(&mut local, backend, full_key)
736 .await
737 .cold_role();
738 let mut flight = Self::new(map, local, role, backend, full_key);
739 flight.lock_id = lock_id;
740 flight
741 }
742 #[cfg(not(all(feature = "reliability", not(target_arch = "wasm32"))))]
743 Self::new(map, local, Role::Leader, backend, full_key)
744 }
745
746 /// Refresh-ahead entry: lead the fill, or `None` to stand down at once.
747 ///
748 /// A refresh runs only while a stale copy is still being served, so it
749 /// never waits for another worker's fill and never re-reads the cache
750 /// (`spec/saas-api.md` API-62/63). It leads only with the lease, or on a
751 /// backend without one. A held in-process flight, a contested lease or a
752 /// failed lock call all mean stand down: the stale copy keeps being
753 /// served, and once it hard-expires the cold-miss path takes over.
754 pub(crate) async fn try_lead(
755 map: &crate::client::SharedFlight,
756 backend: &SharedBackend,
757 full_key: &str,
758 ) -> Option<Self> {
759 let local = map.handle(full_key).try_lock_owned().ok()?;
760 #[cfg_attr(
761 not(all(feature = "reliability", not(target_arch = "wasm32"))),
762 allow(unused_mut)
763 )]
764 let mut flight = Self::new(map, local, Role::Leader, backend, full_key);
765 #[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
766 match LockAttempt::run(&mut flight.local, backend, full_key).await {
767 LockAttempt::Granted(lock_id) => flight.lock_id = Some(lock_id),
768 LockAttempt::Unlockable => {}
769 LockAttempt::Contested { .. } => {
770 // Leave a note: a cold miss queued behind this refresh must
771 // not take the stand-down for a finished fill.
772 flight.local.stand_down = true;
773 return None;
774 }
775 LockAttempt::Failed => return None,
776 }
777 Some(flight)
778 }
779}
780
781/// Free a distributed fill lock; `true` once the backend answered. Best-effort:
782/// a failed unlock is traced (feature `tracing`) and otherwise ignored, since
783/// the lock expires server-side regardless.
784#[cfg(all(feature = "reliability", not(target_arch = "wasm32")))]
785async fn release_lock(backend: &SharedBackend, full_key: &str, lock_id: &str) -> bool {
786 let Some(lockable) = backend.as_lockable() else {
787 return false;
788 };
789 match lockable.release_lock(full_key, lock_id).await {
790 Ok(_) => true,
791 Err(err) => {
792 crate::metrics::trace_lock_error(full_key, "unlock", &err);
793 false
794 }
795 }
796}
797
798#[cfg(all(
799 test,
800 feature = "reliability",
801 feature = "l1",
802 not(feature = "unsync"),
803 not(target_arch = "wasm32")
804))]
805mod tests {
806 use std::sync::atomic::AtomicUsize;
807
808 use async_trait::async_trait;
809
810 use super::*;
811 use crate::backend::{Backend, HealthStatus, LockableBackend};
812 use crate::error::BackendError;
813
814 /// Grants every lock at once and counts unlocks.
815 #[derive(Default)]
816 struct Locks {
817 unlocks: AtomicUsize,
818 }
819
820 #[async_trait]
821 impl Backend for Locks {
822 async fn get(&self, _key: &str) -> Result<Option<Vec<u8>>, BackendError> {
823 Ok(None)
824 }
825 async fn set(
826 &self,
827 _key: &str,
828 _value: Vec<u8>,
829 _ttl: Option<std::time::Duration>,
830 ) -> Result<(), BackendError> {
831 Ok(())
832 }
833 async fn delete(&self, _key: &str) -> Result<bool, BackendError> {
834 Ok(false)
835 }
836 async fn exists(&self, _key: &str) -> Result<bool, BackendError> {
837 Ok(false)
838 }
839 async fn health(&self) -> Result<HealthStatus, BackendError> {
840 Ok(HealthStatus {
841 is_healthy: true,
842 latency_ms: 0.0,
843 backend_type: "locks".to_owned(),
844 details: HashMap::new(),
845 })
846 }
847 fn as_lockable(&self) -> Option<&dyn LockableBackend> {
848 Some(self)
849 }
850 }
851
852 #[async_trait]
853 impl LockableBackend for Locks {
854 async fn acquire_lock(
855 &self,
856 _key: &str,
857 _timeout_ms: u64,
858 ) -> Result<Option<String>, BackendError> {
859 Ok(Some("lease".to_owned()))
860 }
861 async fn release_lock(&self, _key: &str, _lock_id: &str) -> Result<bool, BackendError> {
862 self.unlocks.fetch_add(1, Ordering::SeqCst);
863 Ok(true)
864 }
865 }
866
867 /// A cold-key burst's dead slots go within a lease timeout of the burst,
868 /// though the sweep during it counted them as live.
869 #[tokio::test(start_paused = true)]
870 async fn a_burst_of_dead_slots_is_swept_after_a_lease_timeout() {
871 let map = FlightMap::default();
872 let burst: Vec<_> = (0..200)
873 .map(|i| map.handle(&format!("burst:{i}")))
874 .collect();
875 drop(burst);
876
877 map.handle("early");
878 assert!(
879 map.entries.lock().unwrap().slots.len() > 200,
880 "swept before a lease timeout"
881 );
882
883 tokio::time::advance(std::time::Duration::from_millis(FILL_LOCK_TIMEOUT_MS)).await;
884 map.handle("late");
885 assert!(
886 map.entries.lock().unwrap().slots.len() <= 2,
887 "dead slots retained"
888 );
889 }
890
891 /// While something else holds the key's entry, so it survives the unpin, a
892 /// re-miss does not re-send an unlock the backend already answered.
893 #[tokio::test(start_paused = true)]
894 async fn an_answered_unlock_is_not_resent_while_the_entry_lives() {
895 let locks = Arc::new(Locks::default());
896 let backend: SharedBackend = locks.clone();
897 let map: crate::client::SharedFlight = Arc::default();
898 let _held = map.handle("k");
899
900 SingleFlight::acquire(&map, &backend, "k")
901 .await
902 .__release_detached()
903 .await;
904 while locks.unlocks.load(Ordering::SeqCst) == 0 {
905 tokio::task::yield_now().await;
906 }
907
908 let flight = SingleFlight::acquire(&map, &backend, "k").await;
909 assert_eq!(flight.lock_id.as_deref(), Some("lease"));
910 assert_eq!(
911 locks.unlocks.load(Ordering::SeqCst),
912 1,
913 "answered unlock re-sent"
914 );
915 }
916}