Skip to main content

faucet_core/
state_key.rs

1//! Typed state keys (#882, RFC 0013).
2//!
3//! Every key a pipeline reads or writes in its [`StateStore`] is built here
4//! from its components and rendered one way:
5//!
6//! ```text
7//! [{tenant}::]{pipeline}::{row}[::{sub}][::{shard}][::{suffix}]
8//! [{tenant}::]{pipeline}::{pipeline-level suffix}
9//! ```
10//!
11//! A changed rendering silently looks like a first run — a full re-sync or a
12//! skipped incremental range — so the rendering is versioned
13//! ([`KEY_SCHEME_VERSION`]) and every key can list what earlier schemes would
14//! have rendered for it ([`StateKey::legacy_renderings`]).
15//! [`migrate_state_key`] moves a value stored under one of those keys to the
16//! current key at run start.
17
18use crate::error::FaucetError;
19use crate::state::StateStore;
20
21/// Version of the key rendering this build writes.
22///
23/// * `1` — a child invocation's parent-key value rendered verbatim.
24/// * `2` — the parent-key value is encoded by [`encode_component`] (#840), so a
25///   value holding `::`, spaces or non-ASCII cannot split or break the key.
26pub const KEY_SCHEME_VERSION: u32 = 2;
27
28/// Separator between key components.
29pub const KEY_SEPARATOR: &str = "::";
30
31/// Longest raw value [`encode_component`] keeps verbatim.
32pub const COMPONENT_MAX: usize = 64;
33
34/// What a key holds besides a plain resume position.
35#[derive(Debug, Clone, PartialEq, Eq, Hash)]
36#[non_exhaustive]
37pub enum KeySuffix {
38    /// `::__sla__` — freshness / volume SLA history of an invocation.
39    Sla,
40    /// `::__profiling__` — column-profile baseline of an invocation.
41    Profiling,
42    /// `::__rollback__` — the retained undoable runs of an invocation.
43    RollbackIndex,
44    /// `::__rollback__::{run}` — one undoable run's marker.
45    RollbackRun(String),
46    /// `::__status__` — the last run outcomes of an invocation.
47    Status,
48    /// `::__lease__` — held while a run of the invocation is in flight.
49    Lease,
50    /// `{pipeline}::__replication__` — the snapshot→CDC phase marker.
51    Replication,
52    /// `{pipeline}::__backfill__::{range hash}` — a backfill's progress marker.
53    Backfill(String),
54    /// `{pipeline}::__serve_backfill__::{range hash}` — a served backfill range.
55    ServeBackfill(String),
56}
57
58impl KeySuffix {
59    /// Whether the suffix hangs off the pipeline rather than an invocation.
60    pub fn is_pipeline_level(&self) -> bool {
61        matches!(
62            self,
63            Self::Replication | Self::Backfill(_) | Self::ServeBackfill(_)
64        )
65    }
66
67    fn render_into(&self, out: &mut String) {
68        let (marker, tail) = match self {
69            Self::Sla => ("__sla__", None),
70            Self::Profiling => ("__profiling__", None),
71            Self::RollbackIndex => ("__rollback__", None),
72            Self::RollbackRun(run) => ("__rollback__", Some(run.as_str())),
73            Self::Status => ("__status__", None),
74            Self::Lease => ("__lease__", None),
75            Self::Replication => ("__replication__", None),
76            Self::Backfill(h) => ("__backfill__", Some(h.as_str())),
77            Self::ServeBackfill(h) => ("__serve_backfill__", Some(h.as_str())),
78        };
79        push(out, marker);
80        if let Some(tail) = tail {
81            push(out, tail);
82        }
83    }
84
85    /// The suffix's rendering, without the leading separator.
86    pub fn render(&self) -> String {
87        let mut s = String::new();
88        self.render_into(&mut s);
89        s
90    }
91
92    /// The suffix as it ends a key, separator included (`::__lease__`).
93    pub fn tail(&self) -> String {
94        format!("{KEY_SEPARATOR}{}", self.render())
95    }
96}
97
98fn push(out: &mut String, seg: &str) {
99    if !out.is_empty() {
100        out.push_str(KEY_SEPARATOR);
101    }
102    out.push_str(seg);
103}
104
105#[derive(Debug, Clone, PartialEq, Eq, Hash)]
106enum Sub {
107    /// A child invocation's parent-key value, as the parent produced it.
108    Parent(String),
109    /// A component already in key form (a product tuple, a parsed segment).
110    Encoded(String),
111}
112
113#[derive(Debug, Clone, PartialEq, Eq, Hash)]
114enum Head {
115    Parts {
116        tenant: Option<String>,
117        pipeline: String,
118        row: Option<String>,
119        sub: Option<Sub>,
120    },
121    /// An invocation key [`StateKey::render`] produced earlier.
122    Rendered(String),
123}
124
125/// A state key built from its components. See the [module docs](self).
126#[derive(Debug, Clone, PartialEq, Eq, Hash)]
127pub struct StateKey {
128    head: Head,
129    shard: Option<String>,
130    suffix: Option<KeySuffix>,
131}
132
133impl StateKey {
134    /// The pipeline-level key namespace of `pipeline` (no row yet).
135    pub fn pipeline(pipeline: impl Into<String>) -> Self {
136        Self {
137            head: Head::Parts {
138                tenant: None,
139                pipeline: pipeline.into(),
140                row: None,
141                sub: None,
142            },
143            shard: None,
144            suffix: None,
145        }
146    }
147
148    /// `{pipeline}::{row}` — a matrix row's (or topology sink node's) key.
149    pub fn invocation(pipeline: impl Into<String>, row: impl Into<String>) -> Self {
150        Self::pipeline(pipeline).with_row(row)
151    }
152
153    /// A key derived from an invocation key [`render`](Self::render) already
154    /// produced — the base the per-invocation markers hang off.
155    pub fn from_rendered(base: impl Into<String>) -> Self {
156        Self {
157            head: Head::Rendered(base.into()),
158            shard: None,
159            suffix: None,
160        }
161    }
162
163    /// Place the key under `tenant` (`{tenant}::{pipeline}::…`, #709).
164    pub fn with_tenant(mut self, tenant: Option<impl Into<String>>) -> Self {
165        if let Head::Parts { tenant: t, .. } = &mut self.head {
166            *t = tenant.map(Into::into);
167        }
168        self
169    }
170
171    /// The matrix row / topology node id.
172    pub fn with_row(mut self, id: impl Into<String>) -> Self {
173        if let Head::Parts { row, .. } = &mut self.head {
174            *row = Some(id.into());
175        }
176        self
177    }
178
179    /// A child invocation's parent-key value, raw; rendered through
180    /// [`encode_component`].
181    pub fn with_parent_key(mut self, raw: impl Into<String>) -> Self {
182        if let Head::Parts { sub, .. } = &mut self.head {
183            *sub = Some(Sub::Parent(raw.into()));
184        }
185        self
186    }
187
188    /// A sub-invocation component already in key form (a product tuple suffix).
189    pub fn with_sub(mut self, encoded: impl Into<String>) -> Self {
190        if let Head::Parts { sub, .. } = &mut self.head {
191            *sub = Some(Sub::Encoded(encoded.into()));
192        }
193        self
194    }
195
196    /// A shard of the invocation (`{base}::{shard}`).
197    pub fn with_shard(mut self, shard: impl Into<String>) -> Self {
198        self.shard = Some(shard.into());
199        self
200    }
201
202    /// A marker suffix.
203    pub fn with_suffix(mut self, suffix: KeySuffix) -> Self {
204        self.suffix = Some(suffix);
205        self
206    }
207
208    /// The tenant, when the key is built from components.
209    pub fn tenant(&self) -> Option<&str> {
210        match &self.head {
211            Head::Parts { tenant, .. } => tenant.as_deref(),
212            Head::Rendered(_) => None,
213        }
214    }
215
216    /// The pipeline name, when the key is built from components.
217    pub fn pipeline_name(&self) -> Option<&str> {
218        match &self.head {
219            Head::Parts { pipeline, .. } => Some(pipeline),
220            Head::Rendered(_) => None,
221        }
222    }
223
224    /// The row id, when the key is built from components.
225    pub fn row(&self) -> Option<&str> {
226        match &self.head {
227            Head::Parts { row, .. } => row.as_deref(),
228            Head::Rendered(_) => None,
229        }
230    }
231
232    /// The rendered sub-invocation component.
233    pub fn sub(&self) -> Option<String> {
234        match &self.head {
235            Head::Parts { sub: Some(s), .. } => Some(render_sub(s, KEY_SCHEME_VERSION)),
236            _ => None,
237        }
238    }
239
240    /// The shard id.
241    pub fn shard(&self) -> Option<&str> {
242        self.shard.as_deref()
243    }
244
245    /// The marker suffix.
246    pub fn suffix(&self) -> Option<&KeySuffix> {
247        self.suffix.as_ref()
248    }
249
250    /// `[{tenant}::]{pipeline}` — the first segments of every key of the
251    /// pipeline; for a key built from a rendering, that rendering.
252    pub fn namespace(&self) -> String {
253        match &self.head {
254            Head::Parts {
255                tenant, pipeline, ..
256            } => {
257                let mut s = String::new();
258                if let Some(t) = tenant {
259                    push(&mut s, t);
260                }
261                push(&mut s, pipeline);
262                s
263            }
264            Head::Rendered(base) => base.clone(),
265        }
266    }
267
268    /// The prefix every key under this one starts with (`{key}::`).
269    pub fn children_prefix(&self) -> String {
270        format!("{}{KEY_SEPARATOR}", self.render())
271    }
272
273    /// The key under the current scheme.
274    pub fn render(&self) -> String {
275        self.render_scheme(KEY_SCHEME_VERSION)
276    }
277
278    /// The key as scheme `version` rendered it.
279    pub fn render_scheme(&self, version: u32) -> String {
280        let mut s = match &self.head {
281            Head::Parts { row, sub, .. } => {
282                let mut s = self.namespace();
283                if let Some(row) = row {
284                    push(&mut s, row);
285                }
286                if let Some(sub) = sub {
287                    push(&mut s, &render_sub(sub, version));
288                }
289                s
290            }
291            Head::Rendered(base) => base.clone(),
292        };
293        if let Some(shard) = &self.shard {
294            push(&mut s, shard);
295        }
296        if let Some(suffix) = &self.suffix {
297            suffix.render_into(&mut s);
298        }
299        s
300    }
301
302    /// Every distinct, valid key an earlier scheme rendered for this key, newest
303    /// scheme first — where [`migrate_state_key`] looks when the current key
304    /// holds nothing.
305    pub fn legacy_renderings(&self) -> Vec<String> {
306        let current = self.render();
307        let mut out: Vec<String> = Vec::new();
308        for version in (1..KEY_SCHEME_VERSION).rev() {
309            let old = self.render_scheme(version);
310            if old != current
311                && !out.contains(&old)
312                && crate::state::validate_state_key(&old).is_ok()
313            {
314                out.push(old);
315            }
316        }
317        out
318    }
319
320    /// Decode `key` under the namespace `[{tenant}::]{pipeline}`; `None` when it
321    /// is outside it. A shard id cannot be told from a sub component, so it is
322    /// returned as part of [`sub`](Self::sub); the rendering round-trips.
323    pub fn parse(tenant: Option<&str>, pipeline: &str, key: &str) -> Option<Self> {
324        let ns = StateKey::pipeline(pipeline).with_tenant(tenant).namespace();
325        let rest = key.strip_prefix(&ns)?.strip_prefix(KEY_SEPARATOR)?;
326        if rest.is_empty() {
327            return None;
328        }
329        let segs: Vec<&str> = rest.split(KEY_SEPARATOR).collect();
330        let join = |s: &[&str]| s.join(KEY_SEPARATOR);
331        let base = StateKey::pipeline(pipeline).with_tenant(tenant);
332        match segs[0] {
333            "__replication__" if segs.len() == 1 => {
334                return Some(base.with_suffix(KeySuffix::Replication));
335            }
336            "__backfill__" if segs.len() > 1 => {
337                return Some(base.with_suffix(KeySuffix::Backfill(join(&segs[1..]))));
338            }
339            "__serve_backfill__" if segs.len() > 1 => {
340                return Some(base.with_suffix(KeySuffix::ServeBackfill(join(&segs[1..]))));
341            }
342            _ => {}
343        }
344        let marker = segs.iter().enumerate().skip(1).find_map(|(i, s)| {
345            let suffix = match *s {
346                "__sla__" if i + 1 == segs.len() => KeySuffix::Sla,
347                "__profiling__" if i + 1 == segs.len() => KeySuffix::Profiling,
348                "__status__" if i + 1 == segs.len() => KeySuffix::Status,
349                "__lease__" if i + 1 == segs.len() => KeySuffix::Lease,
350                "__rollback__" if i + 1 < segs.len() => {
351                    KeySuffix::RollbackRun(join(&segs[i + 1..]))
352                }
353                "__rollback__" => KeySuffix::RollbackIndex,
354                _ => return None,
355            };
356            Some((i, suffix))
357        });
358        let end = marker.as_ref().map_or(segs.len(), |(i, _)| *i);
359        // A backfill unit's row id is `backfill::{range hash}::{unit}`.
360        let row_end = if segs[0] == "backfill" && end > 1 {
361            end
362        } else {
363            1
364        };
365        let mut parsed = base.with_row(join(&segs[..row_end]));
366        if end > row_end {
367            parsed = parsed.with_sub(join(&segs[row_end..end]));
368        }
369        if let Some((_, suffix)) = marker {
370            parsed = parsed.with_suffix(suffix);
371        }
372        Some(parsed)
373    }
374}
375
376fn render_sub(sub: &Sub, version: u32) -> String {
377    match sub {
378        Sub::Parent(raw) if version >= 2 => encode_component(raw, true),
379        Sub::Parent(raw) => raw.clone(),
380        Sub::Encoded(s) => s.clone(),
381    }
382}
383
384impl std::fmt::Display for StateKey {
385    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
386        f.write_str(&self.render())
387    }
388}
389
390/// Render one raw value as a state-key component (#789 CLI-52/CLI-53). A value
391/// made only of state-key characters is kept as is; anything else (an email, a
392/// space, non-ASCII, `::`, an over-long value) becomes a readable prefix plus a
393/// stable 64-bit hash of the raw value. `allow_separators` admits `:` and `/`
394/// (a child's parent key, which stands alone after `::`).
395pub fn encode_component(raw: &str, allow_separators: bool) -> String {
396    let plain = |c: char| c.is_ascii_alphanumeric() || matches!(c, '_' | '-' | '.');
397    let legal = |c: char| plain(c) || (allow_separators && matches!(c, ':' | '/'));
398    if !raw.is_empty()
399        && raw.len() <= COMPONENT_MAX
400        && raw.chars().all(legal)
401        && !raw.contains(KEY_SEPARATOR)
402        && !raw.starts_with(':')
403        && !raw.ends_with(':')
404    {
405        return raw.to_owned();
406    }
407    let readable: String = raw.chars().filter(|c| plain(*c)).take(32).collect();
408    let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
409    for b in raw.as_bytes() {
410        hash ^= u64::from(*b);
411        hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
412    }
413    format!("{readable}.h{hash:016x}")
414}
415
416/// The row id of one backfill unit's invocation, `backfill::{range hash}::{unit}`
417/// — namespaced away from every live row's key.
418pub fn backfill_unit_row(range_hash: &str, unit: &str) -> String {
419    ["backfill", range_hash, unit].join(KEY_SEPARATOR)
420}
421
422/// What [`migrate_state_key`] did.
423#[derive(Debug, Clone, PartialEq, Eq)]
424#[non_exhaustive]
425pub enum KeyMigration {
426    /// The current key already holds a value; nothing was looked up.
427    Current,
428    /// No candidate holds a value: a genuine first run.
429    Fresh,
430    /// The value under `from` now lives under the current key.
431    Moved {
432        /// The earlier key the value was found under.
433        from: String,
434    },
435}
436
437/// Move a value an earlier key scheme (or a renamed row) left under one of
438/// `candidates` to `key`, when `key` holds nothing.
439///
440/// The first candidate holding a value wins. The value is claimed with
441/// [`StateStore::compare_and_put`] against an absent entry, so a run that wrote
442/// `key` meanwhile is never overwritten; the old key is deleted afterwards. A
443/// crash in between leaves the value under both keys, and the next run sees the
444/// current key and moves nothing. Counted in
445/// `faucet_state_key_migrations_total{outcome}`.
446pub async fn migrate_state_key(
447    store: &dyn StateStore,
448    key: &str,
449    candidates: &[String],
450) -> Result<KeyMigration, FaucetError> {
451    if candidates.is_empty() {
452        return Ok(KeyMigration::Fresh);
453    }
454    if store.get(key).await?.is_some() {
455        return Ok(KeyMigration::Current);
456    }
457    for old in candidates.iter().filter(|c| c.as_str() != key) {
458        let Some(value) = store.get(old).await? else {
459            continue;
460        };
461        if !store.compare_and_put(key, None, &value).await? {
462            metrics::counter!("faucet_state_key_migrations_total", "outcome" => "raced")
463                .increment(1);
464            return Ok(KeyMigration::Current);
465        }
466        if let Err(e) = store.delete(old).await {
467            tracing::warn!(key = %old, error = %e, "legacy state key could not be removed");
468        }
469        metrics::counter!("faucet_state_key_migrations_total", "outcome" => "moved").increment(1);
470        tracing::info!(from = %old, to = key, "moved state from the key an earlier release used");
471        return Ok(KeyMigration::Moved { from: old.clone() });
472    }
473    Ok(KeyMigration::Fresh)
474}
475
476#[cfg(test)]
477mod tests {
478    use super::*;
479    use crate::state::MemoryStateStore;
480    use serde_json::json;
481
482    #[test]
483    fn renders_every_kind() {
484        let inv = StateKey::invocation("p", "r");
485        assert_eq!(inv.render(), "p::r");
486        assert_eq!(StateKey::pipeline("p").render(), "p");
487        assert_eq!(
488            StateKey::invocation("p", "r")
489                .with_tenant(Some("acme"))
490                .render(),
491            "acme::p::r"
492        );
493        assert_eq!(inv.clone().with_parent_key("42").render(), "p::r::42");
494        assert_eq!(inv.clone().with_sub("a:1/b:2").render(), "p::r::a:1/b:2");
495        assert_eq!(inv.clone().with_shard("s3").render(), "p::r::s3");
496        let cases = [
497            (KeySuffix::Sla, "p::r::__sla__"),
498            (KeySuffix::Profiling, "p::r::__profiling__"),
499            (KeySuffix::RollbackIndex, "p::r::__rollback__"),
500            (KeySuffix::RollbackRun("x".into()), "p::r::__rollback__::x"),
501            (KeySuffix::Status, "p::r::__status__"),
502            (KeySuffix::Lease, "p::r::__lease__"),
503        ];
504        for (suffix, want) in cases {
505            assert!(!suffix.is_pipeline_level());
506            assert_eq!(inv.clone().with_suffix(suffix).render(), want);
507        }
508        for (suffix, want) in [
509            (KeySuffix::Replication, "p::__replication__"),
510            (KeySuffix::Backfill("h".into()), "p::__backfill__::h"),
511            (
512                KeySuffix::ServeBackfill("h".into()),
513                "p::__serve_backfill__::h",
514            ),
515        ] {
516            assert!(suffix.is_pipeline_level());
517            assert_eq!(StateKey::pipeline("p").with_suffix(suffix).render(), want);
518        }
519        assert_eq!(
520            StateKey::from_rendered("p::r::7")
521                .with_suffix(KeySuffix::Status)
522                .to_string(),
523            "p::r::7::__status__"
524        );
525        assert_eq!(
526            KeySuffix::RollbackRun("x".into()).render(),
527            "__rollback__::x"
528        );
529        assert_eq!(inv.children_prefix(), "p::r::");
530        assert_eq!(KeySuffix::Lease.tail(), "::__lease__");
531    }
532
533    #[test]
534    fn accessors_report_components() {
535        let k = StateKey::invocation("p", "r")
536            .with_tenant(Some("t"))
537            .with_parent_key("a b")
538            .with_shard("s")
539            .with_suffix(KeySuffix::Lease);
540        assert_eq!(k.tenant(), Some("t"));
541        assert_eq!(k.pipeline_name(), Some("p"));
542        assert_eq!(k.row(), Some("r"));
543        assert_eq!(k.sub(), Some(encode_component("a b", true)));
544        assert_eq!(k.shard(), Some("s"));
545        assert_eq!(k.suffix(), Some(&KeySuffix::Lease));
546        assert_eq!(k.namespace(), "t::p");
547        let r = StateKey::from_rendered("x::y")
548            .with_tenant(Some("ignored"))
549            .with_row("ignored")
550            .with_parent_key("ignored")
551            .with_sub("ignored");
552        assert_eq!(r.render(), "x::y");
553        assert_eq!(
554            (r.tenant(), r.pipeline_name(), r.row(), r.sub()),
555            (None, None, None, None)
556        );
557        assert_eq!(r.namespace(), "x::y");
558        assert_eq!(StateKey::invocation("p", "r").sub(), None);
559    }
560
561    #[test]
562    fn parent_keys_are_encoded_and_remember_the_verbatim_scheme() {
563        let k = StateKey::invocation("p", "r").with_parent_key("a b");
564        assert_ne!(k.render(), "p::r::a b");
565        assert_eq!(k.render_scheme(1), "p::r::a b");
566        // `a b` is not a valid key, so scheme 1 has nothing to offer.
567        assert!(k.legacy_renderings().is_empty());
568        let k = StateKey::invocation("p", "r").with_parent_key("x::y");
569        assert_eq!(k.legacy_renderings(), vec!["p::r::x::y".to_string()]);
570        let k = StateKey::invocation("p", "r").with_parent_key("plain");
571        assert!(k.legacy_renderings().is_empty());
572    }
573
574    #[test]
575    fn encode_component_keeps_legal_values_and_hashes_the_rest() {
576        assert_eq!(encode_component("cust-42_a.b", false), "cust-42_a.b");
577        assert_eq!(encode_component("a/b:c", true), "a/b:c");
578        assert_ne!(encode_component("a/b:c", false), "a/b:c");
579        for raw in ["", "x::y", ":x", "y:", "émoji", &"z".repeat(65)] {
580            let c = encode_component(raw, true);
581            assert!(c.contains(".h"), "{raw:?} → {c}");
582            crate::state::validate_state_key(&format!("p::r::{c}")).unwrap();
583        }
584    }
585
586    #[test]
587    fn parse_round_trips_every_kind() {
588        let keys = [
589            "p::r",
590            "p::r::42",
591            "p::r::a:1/b:2",
592            "p::r::__sla__",
593            "p::r::__profiling__",
594            "p::r::__rollback__",
595            "p::r::__rollback__::run1",
596            "p::r::7::__status__",
597            "p::r::__lease__",
598            "p::__replication__",
599            "p::__backfill__::abc",
600            "p::__serve_backfill__::abc",
601            "p::backfill::h::2026-01-01",
602            "p::backfill::h::u1::__status__",
603            "p::backfill",
604        ];
605        for key in keys {
606            let k = StateKey::parse(None, "p", key).unwrap_or_else(|| panic!("{key}"));
607            assert_eq!(k.render(), key);
608        }
609        assert_eq!(backfill_unit_row("h", "u"), "backfill::h::u");
610        let k = StateKey::parse(None, "p", "p::backfill::h::u1::__lease__").unwrap();
611        assert_eq!(k.row(), Some("backfill::h::u1"));
612        assert_eq!(k.suffix(), Some(&KeySuffix::Lease));
613        let k = StateKey::parse(Some("t"), "p", "t::p::r::9").unwrap();
614        assert_eq!(
615            (k.tenant(), k.row(), k.sub()),
616            (Some("t"), Some("r"), Some("9".into()))
617        );
618        assert!(StateKey::parse(None, "p", "q::r").is_none());
619        assert!(StateKey::parse(None, "p", "p::").is_none());
620        assert!(StateKey::parse(None, "p", "px::r").is_none());
621    }
622
623    #[tokio::test]
624    async fn migration_moves_once_and_never_clobbers() {
625        let store = MemoryStateStore::new();
626        assert_eq!(
627            migrate_state_key(&store, "p::new", &[]).await.unwrap(),
628            KeyMigration::Fresh
629        );
630        assert_eq!(
631            migrate_state_key(&store, "p::new", &["p::old".into()])
632                .await
633                .unwrap(),
634            KeyMigration::Fresh
635        );
636        store.put("p::old", &json!({"page": 3})).await.unwrap();
637        assert_eq!(
638            migrate_state_key(&store, "p::new", &["p::new".into(), "p::old".into()])
639                .await
640                .unwrap(),
641            KeyMigration::Moved {
642                from: "p::old".into()
643            }
644        );
645        assert_eq!(store.get("p::new").await.unwrap(), Some(json!({"page": 3})));
646        assert_eq!(store.get("p::old").await.unwrap(), None);
647        store.put("p::old", &json!(1)).await.unwrap();
648        assert_eq!(
649            migrate_state_key(&store, "p::new", &["p::old".into()])
650                .await
651                .unwrap(),
652            KeyMigration::Current
653        );
654        assert_eq!(store.get("p::new").await.unwrap(), Some(json!({"page": 3})));
655    }
656
657    struct Racing(MemoryStateStore);
658
659    #[async_trait::async_trait]
660    impl StateStore for Racing {
661        async fn get(&self, key: &str) -> Result<Option<serde_json::Value>, FaucetError> {
662            self.0.get(key).await
663        }
664        async fn put(&self, key: &str, value: &serde_json::Value) -> Result<(), FaucetError> {
665            self.0.put(key, value).await
666        }
667        async fn delete(&self, key: &str) -> Result<(), FaucetError> {
668            self.0.delete(key).await
669        }
670        async fn compare_and_put(
671            &self,
672            _key: &str,
673            _expected: Option<&serde_json::Value>,
674            _value: &serde_json::Value,
675        ) -> Result<bool, FaucetError> {
676            Ok(false)
677        }
678    }
679
680    /// Removes nothing: every delete fails.
681    struct Sticky(MemoryStateStore);
682
683    #[async_trait::async_trait]
684    impl StateStore for Sticky {
685        async fn get(&self, key: &str) -> Result<Option<serde_json::Value>, FaucetError> {
686            self.0.get(key).await
687        }
688        async fn put(&self, key: &str, value: &serde_json::Value) -> Result<(), FaucetError> {
689            self.0.put(key, value).await
690        }
691        async fn delete(&self, _key: &str) -> Result<(), FaucetError> {
692            Err(FaucetError::State("read-only".into()))
693        }
694    }
695
696    /// A legacy key that cannot be removed still moves: the value is under the
697    /// current key, and the stale copy is left behind (a later run sees the
698    /// current key first).
699    #[tokio::test]
700    async fn a_legacy_key_that_cannot_be_removed_still_moves() {
701        let store = Sticky(MemoryStateStore::new());
702        store.put("p::old", &json!(2)).await.unwrap();
703        migrate_state_key(&store, "p::new", &["p::old".into()])
704            .await
705            .unwrap();
706        assert_eq!(store.get("p::new").await.unwrap(), Some(json!(2)));
707        assert_eq!(store.get("p::old").await.unwrap(), Some(json!(2)));
708        let racing = Racing(MemoryStateStore::new());
709        racing.put("k", &json!(1)).await.unwrap();
710        racing.delete("k").await.unwrap();
711        assert_eq!(racing.get("k").await.unwrap(), None);
712    }
713
714    #[tokio::test]
715    async fn a_concurrent_writer_wins_the_race() {
716        let store = Racing(MemoryStateStore::new());
717        store.put("p::old", &json!(1)).await.unwrap();
718        assert_eq!(
719            migrate_state_key(&store, "p::new", &["p::old".into()])
720                .await
721                .unwrap(),
722            KeyMigration::Current
723        );
724        assert_eq!(store.get("p::old").await.unwrap(), Some(json!(1)));
725    }
726}