faucet_core/shard.rs
1//! Source sharding for clustered (Mode B) execution.
2//!
3//! A *shardable* source can split its work into independent [`ShardSpec`]s that
4//! different cluster workers process concurrently — so one logical pipeline over
5//! a single large source (a big S3 prefix, a billion-row table) scales
6//! horizontally instead of being capped at one worker's throughput.
7//!
8//! The capability is exposed as **defaulted methods on the
9//! [`Source`](crate::Source) trait**
10//! ([`enumerate_shards`](crate::Source::enumerate_shards),
11//! [`apply_shard`](crate::Source::apply_shard),
12//! [`is_shardable`](crate::Source::is_shardable)) rather than a separate trait,
13//! so it composes with the existing `Box<dyn Source>` object model with no
14//! downcast and stays fully backward compatible: a source that does not override
15//! them reports `is_shardable() == false` and enumerates to a single
16//! whole-dataset shard, i.e. it behaves exactly as today.
17//!
18//! ## Lifecycle
19//!
20//! 1. A coordinator calls [`enumerate_shards`](crate::Source::enumerate_shards)
21//! once per run (idempotently) to discover the shard set.
22//! 2. Each shard is persisted and claimed by exactly one worker.
23//! 3. The claiming worker calls [`apply_shard`](crate::Source::apply_shard) to
24//! narrow its source instance to that one shard before streaming.
25//! 4. Each shard carries its own state-key suffix so resume is per-shard: a
26//! reassigned shard resumes where its dead owner left off.
27
28use serde::{Deserialize, Serialize};
29
30/// One independently-processable slice of a shardable source.
31///
32/// The [`descriptor`](ShardSpec::descriptor) is opaque to the coordinator — only
33/// the producing/consuming connector interprets it (e.g. `{"pk_range":[0,1000]}`
34/// for a key-range split, `{"prefix":"dt=2026-06-11/"}` for an object split). It
35/// must round-trip through JSON unchanged so the coordinator can persist it and
36/// hand it back to [`apply_shard`](crate::Source::apply_shard) on the worker that
37/// claims the shard.
38#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
39pub struct ShardSpec {
40 /// Stable identifier, unique within a run's shard set. Also used as the
41 /// per-shard state-key suffix (`{run}::{shard.id}`), so it must be a valid
42 /// state-key segment (no `:` — see
43 /// [`validate_state_key`](crate::state::validate_state_key)).
44 pub id: String,
45
46 /// Connector-specific, opaque shard descriptor. Persisted verbatim and
47 /// passed back to [`apply_shard`](crate::Source::apply_shard).
48 pub descriptor: serde_json::Value,
49
50 /// Optional relative size estimate (rows / bytes / objects — the unit is the
51 /// connector's choice, only comparisons within one run's set matter) used
52 /// for skew-aware assignment. `None` when the source cannot cheaply
53 /// estimate; the coordinator then falls back to count-based balancing.
54 #[serde(default, skip_serializing_if = "Option::is_none")]
55 pub size_estimate: Option<u64>,
56}
57
58impl ShardSpec {
59 /// The single whole-dataset shard a non-shardable source enumerates to.
60 ///
61 /// Its [`id`](ShardSpec::id) is `"0"` and its descriptor is JSON `null`;
62 /// [`apply_shard`](crate::Source::apply_shard)'s default ignores it, so the
63 /// source streams its entire dataset unchanged.
64 pub fn whole() -> Self {
65 Self {
66 id: "0".to_string(),
67 descriptor: serde_json::Value::Null,
68 size_estimate: None,
69 }
70 }
71
72 /// A shard with the given id and descriptor and no size estimate.
73 pub fn new(id: impl Into<String>, descriptor: serde_json::Value) -> Self {
74 Self {
75 id: id.into(),
76 descriptor,
77 size_estimate: None,
78 }
79 }
80
81 /// Attach a relative size estimate (builder style).
82 pub fn with_size(mut self, size: u64) -> Self {
83 self.size_estimate = Some(size);
84 self
85 }
86
87 /// Whether this is the whole-dataset shard (id `"0"`, null descriptor) — the
88 /// no-op shard a non-shardable source produces.
89 pub fn is_whole(&self) -> bool {
90 self.id == "0" && self.descriptor.is_null()
91 }
92}
93
94#[cfg(test)]
95mod tests {
96 use super::*;
97 use serde_json::json;
98
99 #[test]
100 fn whole_is_the_no_op_shard() {
101 let w = ShardSpec::whole();
102 assert_eq!(w.id, "0");
103 assert!(w.descriptor.is_null());
104 assert_eq!(w.size_estimate, None);
105 assert!(w.is_whole());
106 }
107
108 #[test]
109 fn new_and_with_size() {
110 let s = ShardSpec::new("3", json!({ "partition": 3 })).with_size(42);
111 assert_eq!(s.id, "3");
112 assert_eq!(s.descriptor, json!({ "partition": 3 }));
113 assert_eq!(s.size_estimate, Some(42));
114 assert!(!s.is_whole(), "a real shard is not the whole-dataset shard");
115 }
116
117 #[test]
118 fn round_trips_through_json_with_size_omitted_when_none() {
119 let s = ShardSpec::new("a", json!({ "prefix": "dt=2026-06-11/" }));
120 let text = serde_json::to_string(&s).unwrap();
121 // size_estimate is omitted from the wire form when None.
122 assert!(
123 !text.contains("size_estimate"),
124 "None size is skipped: {text}"
125 );
126 let back: ShardSpec = serde_json::from_str(&text).unwrap();
127 assert_eq!(back, s);
128 }
129
130 #[test]
131 fn round_trips_with_size_present() {
132 let s = ShardSpec::new("b", json!({ "pk_range": [0, 1000] })).with_size(1000);
133 let back: ShardSpec = serde_json::from_str(&serde_json::to_string(&s).unwrap()).unwrap();
134 assert_eq!(back, s);
135 assert_eq!(back.size_estimate, Some(1000));
136 }
137
138 // A non-whole shard with id "0" but a non-null descriptor is NOT treated as
139 // the whole-dataset shard (the descriptor disambiguates).
140 #[test]
141 fn id_zero_with_descriptor_is_not_whole() {
142 let s = ShardSpec::new("0", json!({ "partition": 0 }));
143 assert!(!s.is_whole());
144 }
145}