1use crate::error::FaucetError;
19use crate::state::StateStore;
20
21pub const KEY_SCHEME_VERSION: u32 = 2;
27
28pub const KEY_SEPARATOR: &str = "::";
30
31pub const COMPONENT_MAX: usize = 64;
33
34#[derive(Debug, Clone, PartialEq, Eq, Hash)]
36#[non_exhaustive]
37pub enum KeySuffix {
38 Sla,
40 Profiling,
42 RollbackIndex,
44 RollbackRun(String),
46 Status,
48 Lease,
50 Replication,
52 Backfill(String),
54 ServeBackfill(String),
56}
57
58impl KeySuffix {
59 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 pub fn render(&self) -> String {
87 let mut s = String::new();
88 self.render_into(&mut s);
89 s
90 }
91
92 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 Parent(String),
109 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 Rendered(String),
123}
124
125#[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 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 pub fn invocation(pipeline: impl Into<String>, row: impl Into<String>) -> Self {
150 Self::pipeline(pipeline).with_row(row)
151 }
152
153 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 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 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 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 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 pub fn with_shard(mut self, shard: impl Into<String>) -> Self {
198 self.shard = Some(shard.into());
199 self
200 }
201
202 pub fn with_suffix(mut self, suffix: KeySuffix) -> Self {
204 self.suffix = Some(suffix);
205 self
206 }
207
208 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 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 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 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 pub fn shard(&self) -> Option<&str> {
242 self.shard.as_deref()
243 }
244
245 pub fn suffix(&self) -> Option<&KeySuffix> {
247 self.suffix.as_ref()
248 }
249
250 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 pub fn children_prefix(&self) -> String {
270 format!("{}{KEY_SEPARATOR}", self.render())
271 }
272
273 pub fn render(&self) -> String {
275 self.render_scheme(KEY_SCHEME_VERSION)
276 }
277
278 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 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 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 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
390pub 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
416pub fn backfill_unit_row(range_hash: &str, unit: &str) -> String {
419 ["backfill", range_hash, unit].join(KEY_SEPARATOR)
420}
421
422#[derive(Debug, Clone, PartialEq, Eq)]
424#[non_exhaustive]
425pub enum KeyMigration {
426 Current,
428 Fresh,
430 Moved {
432 from: String,
434 },
435}
436
437pub 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 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 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 #[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}