ifc_lite_server/
in_flight.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4
5//! In-memory marker for "a data-model write is running right now for this
6//! cache key" (issue #5129).
7//!
8//! The key is the data-model ENTRY key (`data_model_cache_key`), not the
9//! request key it is built from (#6034): the full and the rooted-only tables
10//! are separate entries, and a fill of one says nothing about the other.
11//!
12//! `get_data_model` needs to tell a client "not yet, keep polling" apart from
13//! "not for this key, stop asking": the first is 202, the second 404. A cache
14//! miss alone cannot make that distinction — the entry could be absent because
15//! nothing ever asked for it, which is exactly the case the issue's 404 is
16//! for. This map is the missing signal: a key is present in it for the
17//! lifetime of every background task that could still write it
18//! (`parse_parquet_stream`'s data-model fill), and gone only once every such
19//! task for that key has ended, however each one ends.
20//!
21//! Refcounted, not a plain set (PR #5134 review): two concurrent uploads of
22//! the same content each spawn their own fill and each call `begin` with the
23//! same cache key. A plain set's `remove` on either `Drop` would clear the
24//! key while the other fill is still running, and a `GET` racing that window
25//! would see 404 for a fill that is, in fact, still in progress.
26
27use rustc_hash::FxHashMap;
28use std::sync::{Arc, Mutex};
29
30/// The refcounted set of cache keys with a data-model write in flight. One
31/// instance lives on `AppState`, shared by every request. The count is how
32/// many live guards currently name that key, not a byte size or anything
33/// meaningful beyond "still positive".
34#[derive(Default)]
35pub struct InFlightKeys(Mutex<FxHashMap<String, usize>>);
36
37impl InFlightKeys {
38    /// Mark `key` as in flight and return a guard that un-marks it on drop.
39    ///
40    /// Takes `self` as `Arc` so the returned guard can outlive the borrow —
41    /// it is moved into the spawned task, which outlives the request handler
42    /// that calls `begin`.
43    pub fn begin(self: &Arc<Self>, key: String) -> InFlightGuard {
44        if let Ok(mut keys) = self.0.lock() {
45            *keys.entry(key.clone()).or_insert(0) += 1;
46        }
47        InFlightGuard {
48            keys: self.clone(),
49            key,
50        }
51    }
52
53    /// Whether a write for `key` is in flight right now.
54    ///
55    /// A lock-poisoned map (a prior holder panicked mid-mutation) answers
56    /// `false`: the safe direction is a 404 telling the client to re-upload,
57    /// not a 202 that polls a key nothing can ever finish writing.
58    pub fn contains(&self, key: &str) -> bool {
59        self.0
60            .lock()
61            .map(|keys| keys.contains_key(key))
62            .unwrap_or(false)
63    }
64}
65
66/// RAII marker returned by [`InFlightKeys::begin`]. Decrements its key's
67/// count on drop -- on task completion, on early return (e.g. admission
68/// saturated), and on panic unwind alike -- removing the key only once the
69/// count reaches zero, so one of two concurrent fills for the same key ending
70/// can never un-mark the other's still-running fill.
71pub struct InFlightGuard {
72    keys: Arc<InFlightKeys>,
73    key: String,
74}
75
76impl Drop for InFlightGuard {
77    fn drop(&mut self) {
78        if let Ok(mut keys) = self.keys.0.lock() {
79            if let Some(count) = keys.get_mut(&self.key) {
80                *count -= 1;
81                if *count == 0 {
82                    keys.remove(&self.key);
83                }
84            }
85        }
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92
93    /// `contains` is true only between `begin` and the guard's drop.
94    #[test]
95    fn contains_true_only_while_guard_is_alive() {
96        let keys = Arc::new(InFlightKeys::default());
97        assert!(!keys.contains("k"));
98        let guard = keys.begin("k".to_string());
99        assert!(keys.contains("k"));
100        drop(guard);
101        assert!(!keys.contains("k"));
102    }
103
104    /// Two different keys don't interfere with each other.
105    #[test]
106    fn keys_are_independent() {
107        let keys = Arc::new(InFlightKeys::default());
108        let guard_a = keys.begin("a".to_string());
109        assert!(keys.contains("a"));
110        assert!(!keys.contains("b"));
111        let guard_b = keys.begin("b".to_string());
112        assert!(keys.contains("b"));
113        drop(guard_a);
114        assert!(!keys.contains("a"));
115        assert!(keys.contains("b"));
116        drop(guard_b);
117        assert!(!keys.contains("b"));
118    }
119
120    /// Two concurrent fills for the SAME key (issue #5134 review: two
121    /// uploads of identical content each spawning their own data-model
122    /// task): dropping one guard must not un-mark the other's still-running
123    /// fill. Only the last drop clears the key.
124    #[test]
125    fn two_guards_for_the_same_key_require_both_to_drop() {
126        let keys = Arc::new(InFlightKeys::default());
127        let first = keys.begin("k".to_string());
128        let second = keys.begin("k".to_string());
129        assert!(keys.contains("k"));
130
131        drop(first);
132        assert!(
133            keys.contains("k"),
134            "the second fill is still running; the key must still read as in-flight"
135        );
136
137        drop(second);
138        assert!(!keys.contains("k"), "the last guard's drop must clear the key");
139    }
140}
ifc_lite_server/in_flight.rs

ifc_lite_server/
in_flight.rs
