spate_test/
coordination.rs

1//! Scripted coordination for testing coordinated sources.
2//!
3//! [`scripted_coordinator`] pairs a [`SplitCoordinator`] implementation with a
4//! [`CoordinatorScript`] handle. The test scripts ownership events and commit
5//! outcomes, the source under test runs its real driver choreography, and the
6//! script observes every commit, failure report, and release. There is no
7//! store and no clock, and behavior is deterministic.
8
9#![expect(deprecated, reason = "v0.4.0 removes this with spate-coordination")]
10
11use spate_core::coordination::ControlWaker;
12use spate_core::coordination::{
13    CoordinationError, CoordinationErrorKind, CoordinationEvent, LeaseEpoch, SplitCoordinator,
14    SplitId, SplitPlanner, SplitProgress, SplitSpec,
15};
16use std::collections::{HashMap, VecDeque};
17use std::sync::{Arc, Mutex};
18
19#[derive(Default)]
20struct State {
21    events: Vec<CoordinationEvent>,
22    commit_outcomes: HashMap<SplitId, VecDeque<CoordinationErrorKind>>,
23    commits: Vec<(SplitId, SplitProgress)>,
24    final_commits: Vec<Vec<(SplitId, SplitProgress)>>,
25    fail_outcomes: HashMap<SplitId, VecDeque<CoordinationErrorKind>>,
26    failed: Vec<(SplitId, LeaseEpoch, String)>,
27    released: Vec<SplitId>,
28    departed: bool,
29    planner: Option<Box<dyn SplitPlanner>>,
30    started: bool,
31    waker: Option<ControlWaker>,
32}
33
34/// A [`SplitCoordinator`] whose events and outcomes are scripted by the
35/// paired [`CoordinatorScript`]. Build both with [`scripted_coordinator`].
36#[derive(Debug)]
37#[deprecated(
38    since = "0.3.0",
39    note = "removed in v0.4.0, where a work feed replaces split coordination; see https://github.com/spate-etl/spate/issues/1088"
40)]
41pub struct ScriptedCoordinator {
42    state: Arc<Mutex<State>>,
43}
44
45/// Scripting and observation handle for a [`ScriptedCoordinator`].
46///
47/// Events queued between two `poll` calls are delivered as **one batch**,
48/// matching the trait contract ("all pending events at once"). Queue a
49/// [`lose`](CoordinatorScript::lose) and a [`gain`](CoordinatorScript::gain)
50/// back to back to exercise same-batch interleavings.
51#[derive(Clone, Debug)]
52#[deprecated(
53    since = "0.3.0",
54    note = "removed in v0.4.0, where a work feed replaces split coordination; see https://github.com/spate-etl/spate/issues/1088"
55)]
56pub struct CoordinatorScript {
57    state: Arc<Mutex<State>>,
58}
59
60impl std::fmt::Debug for State {
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62        f.debug_struct("State")
63            .field("pending_events", &self.events.len())
64            .field("commits", &self.commits.len())
65            .field("started", &self.started)
66            .finish_non_exhaustive()
67    }
68}
69
70/// A scripted coordinator and its handle.
71#[must_use]
72#[deprecated(
73    since = "0.3.0",
74    note = "removed in v0.4.0, where a work feed replaces split coordination; see https://github.com/spate-etl/spate/issues/1088"
75)]
76pub fn scripted_coordinator() -> (ScriptedCoordinator, CoordinatorScript) {
77    let state = Arc::new(Mutex::new(State::default()));
78    (
79        ScriptedCoordinator {
80            state: Arc::clone(&state),
81        },
82        CoordinatorScript { state },
83    )
84}
85
86impl CoordinatorScript {
87    /// Queue a [`CoordinationEvent::Gained`] for the next poll.
88    pub fn gain(&self, split: SplitSpec, epoch: u64, progress: Option<SplitProgress>) {
89        self.lock().events.push(CoordinationEvent::Gained {
90            split,
91            epoch: LeaseEpoch(epoch),
92            progress,
93        });
94        self.wake();
95    }
96
97    /// Queue a [`CoordinationEvent::Lost`] for the next poll.
98    pub fn lose(&self, split: &SplitId) {
99        self.lock().events.push(CoordinationEvent::Lost {
100            split: split.clone(),
101        });
102        self.wake();
103    }
104
105    /// Queue a [`CoordinationEvent::Quarantined`] for the next poll.
106    pub fn quarantine(&self, split: &SplitId, attempts: u32) {
107        self.lock().events.push(CoordinationEvent::Quarantined {
108            split: split.clone(),
109            attempts,
110        });
111        self.wake();
112    }
113
114    /// Queue [`CoordinationEvent::AllComplete`] for the next poll.
115    pub fn all_complete(&self) {
116        self.lock().events.push(CoordinationEvent::AllComplete);
117        self.wake();
118    }
119
120    /// Queue [`CoordinationEvent::Stalled`] for the next poll.
121    pub fn stalled(&self, completed: u64, quarantined: u64) {
122        self.lock().events.push(CoordinationEvent::Stalled {
123            completed,
124            quarantined,
125        });
126        self.wake();
127    }
128
129    /// Script the outcome of the next `commit` for `split` (repeat to
130    /// script a sequence). Unscripted commits succeed and are recorded.
131    pub fn fail_next_commit(&self, split: &SplitId, kind: CoordinationErrorKind) {
132        self.lock()
133            .commit_outcomes
134            .entry(split.clone())
135            .or_default()
136            .push_back(kind);
137    }
138
139    /// Every successful commit so far, in order.
140    #[must_use]
141    pub fn commits(&self) -> Vec<(SplitId, SplitProgress)> {
142        self.lock().commits.clone()
143    }
144
145    /// Each final commit batch the source sent, in order.
146    #[must_use]
147    pub fn final_commits(&self) -> Vec<Vec<(SplitId, SplitProgress)>> {
148        self.lock().final_commits.clone()
149    }
150
151    /// The last successful commit for `split`, if any.
152    #[must_use]
153    pub fn last_commit(&self, split: &SplitId) -> Option<SplitProgress> {
154        self.lock()
155            .commits
156            .iter()
157            .rev()
158            .find(|(s, _)| s == split)
159            .map(|(_, p)| p.clone())
160    }
161
162    /// Script the outcome of the next `fail` for `split` (repeat to script a
163    /// sequence). Unscripted reports succeed. Every attempt is recorded by
164    /// [`failed`](Self::failed), refused or not.
165    pub fn fail_next_report(&self, split: &SplitId, kind: CoordinationErrorKind) {
166        self.lock()
167            .fail_outcomes
168            .entry(split.clone())
169            .or_default()
170            .push_back(kind);
171    }
172
173    /// Every `fail` report attempt so far, in order, with the epoch of the
174    /// tenancy it reported.
175    #[must_use]
176    pub fn failed(&self) -> Vec<(SplitId, LeaseEpoch, String)> {
177        self.lock().failed.clone()
178    }
179
180    /// Every released split so far, in release order.
181    #[must_use]
182    pub fn released(&self) -> Vec<SplitId> {
183        self.lock().released.clone()
184    }
185
186    /// Whether `depart` ran. Its held splits are also in
187    /// [`released`](CoordinatorScript::released).
188    #[must_use]
189    pub fn departed(&self) -> bool {
190        self.lock().departed
191    }
192
193    /// Whether `start` ran.
194    #[must_use]
195    pub fn started(&self) -> bool {
196        self.lock().started
197    }
198
199    /// Take the planner captured at `start`. Run it directly to inspect its
200    /// plan, then feed the resulting splits back through
201    /// [`gain`](CoordinatorScript::gain) to test a source's planner and its
202    /// driver choreography together.
203    #[must_use]
204    pub fn take_planner(&self) -> Option<Box<dyn SplitPlanner>> {
205        self.lock().planner.take()
206    }
207
208    fn lock(&self) -> std::sync::MutexGuard<'_, State> {
209        self.state.lock().expect("coordinator script poisoned")
210    }
211
212    /// Wake the driver's control-plane park so a scripted event is picked
213    /// up immediately instead of waiting out the caller's poll timeout.
214    fn wake(&self) {
215        if let Some(w) = &self.lock().waker {
216            w.wake();
217        }
218    }
219}
220
221impl SplitCoordinator for ScriptedCoordinator {
222    fn start(&mut self, planner: Box<dyn SplitPlanner>) -> Result<(), CoordinationError> {
223        let mut state = self.state.lock().expect("coordinator script poisoned");
224        state.planner = Some(planner);
225        state.started = true;
226        Ok(())
227    }
228
229    fn set_waker(&mut self, waker: ControlWaker) {
230        self.state
231            .lock()
232            .expect("coordinator script poisoned")
233            .waker = Some(waker);
234    }
235
236    fn poll(&mut self) -> Result<Vec<CoordinationEvent>, CoordinationError> {
237        // Deterministic and non-blocking. Queued events are one batch, and an
238        // empty queue returns immediately rather than waiting out the timeout.
239        Ok(std::mem::take(
240            &mut self
241                .state
242                .lock()
243                .expect("coordinator script poisoned")
244                .events,
245        ))
246    }
247
248    fn commit(
249        &mut self,
250        split: &SplitId,
251        progress: &SplitProgress,
252    ) -> Result<(), CoordinationError> {
253        let mut state = self.state.lock().expect("coordinator script poisoned");
254        if let Some(kinds) = state.commit_outcomes.get_mut(split)
255            && let Some(kind) = kinds.pop_front()
256        {
257            return Err(CoordinationError::new(
258                kind,
259                format!("scripted {kind:?} for split {split}"),
260            ));
261        }
262        state.commits.push((split.clone(), progress.clone()));
263        Ok(())
264    }
265
266    fn commit_final(
267        &mut self,
268        commits: &[(SplitId, SplitProgress)],
269    ) -> Vec<Result<(), CoordinationError>> {
270        self.state
271            .lock()
272            .expect("coordinator script poisoned")
273            .final_commits
274            .push(commits.to_vec());
275        commits
276            .iter()
277            .map(|(split, progress)| self.commit(split, progress))
278            .collect()
279    }
280
281    fn fail(
282        &mut self,
283        split: &SplitId,
284        epoch: LeaseEpoch,
285        reason: &str,
286    ) -> Result<(), CoordinationError> {
287        let mut state = self.state.lock().expect("coordinator script poisoned");
288        state
289            .failed
290            .push((split.clone(), epoch, reason.to_string()));
291        if let Some(kinds) = state.fail_outcomes.get_mut(split)
292            && let Some(kind) = kinds.pop_front()
293        {
294            return Err(CoordinationError::new(
295                kind,
296                format!("scripted {kind:?} for split {split}"),
297            ));
298        }
299        Ok(())
300    }
301
302    fn release(&mut self, splits: &[SplitId]) -> Result<(), CoordinationError> {
303        self.state
304            .lock()
305            .expect("coordinator script poisoned")
306            .released
307            .extend(splits.iter().cloned());
308        Ok(())
309    }
310
311    fn depart(&mut self, held: &[SplitId]) -> Result<(), CoordinationError> {
312        let mut state = self.state.lock().expect("coordinator script poisoned");
313        state.departed = true;
314        state.released.extend(held.iter().cloned());
315        Ok(())
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322    use spate_core::coordination::{PlanContext, PlanFinality, SplitPlan};
323
324    struct OneSplit;
325
326    impl SplitPlanner for OneSplit {
327        fn fingerprint(&self) -> String {
328            "test:v1".into()
329        }
330
331        fn plan(&mut self, _ctx: PlanContext<'_>) -> Result<SplitPlan, CoordinationError> {
332            Ok(SplitPlan::new(
333                vec![spate_core::coordination::PlannedSplit::new(SplitSpec::new(
334                    SplitId::new("only").unwrap(),
335                    b"all of it".to_vec(),
336                ))],
337                PlanFinality::Final,
338            ))
339        }
340    }
341
342    #[test]
343    fn scripts_events_and_observes_interactions() {
344        let (mut coordinator, script) = scripted_coordinator();
345        coordinator.start(Box::new(OneSplit)).unwrap();
346        assert!(script.started());
347
348        // The captured planner is runnable by the test.
349        let mut planner = script.take_planner().expect("planner captured");
350        let plan = planner.plan(PlanContext::new(None, 1)).unwrap();
351        assert_eq!(plan.splits.len(), 1);
352        let split = plan.splits[0].spec.clone();
353        let id = split.id.clone();
354
355        // lose+gain queued together arrive as one batch, in order.
356        script.gain(split, 1, None);
357        script.lose(&id);
358        let batch = coordinator.poll().unwrap();
359        assert_eq!(batch.len(), 2);
360        assert!(matches!(batch[0], CoordinationEvent::Gained { .. }));
361        assert!(matches!(batch[1], CoordinationEvent::Lost { .. }));
362        assert!(coordinator.poll().unwrap().is_empty());
363
364        // Scripted commit outcomes drain in order, then commits succeed.
365        script.fail_next_commit(&id, CoordinationErrorKind::Retryable);
366        let progress = SplitProgress::new(5, vec![]);
367        let err = coordinator.commit(&id, &progress).unwrap_err();
368        assert_eq!(err.kind, CoordinationErrorKind::Retryable);
369        coordinator.commit(&id, &progress).unwrap();
370        assert_eq!(script.commits().len(), 1);
371        assert_eq!(script.last_commit(&id).unwrap().watermark, 5);
372
373        coordinator.fail(&id, LeaseEpoch(1), "poison").unwrap();
374        assert_eq!(script.failed()[0].1, LeaseEpoch(1));
375        assert_eq!(script.failed()[0].2, "poison");
376        coordinator.release(std::slice::from_ref(&id)).unwrap();
377        assert_eq!(script.released(), vec![id]);
378        assert!(!script.departed());
379        coordinator.depart(&[]).unwrap();
380        assert!(script.departed());
381    }
382
383    #[test]
384    fn final_commits_records_each_batch_and_applies_scripted_outcomes() {
385        let (mut coordinator, script) = scripted_coordinator();
386        let ids: Vec<SplitId> = ["a", "b", "c"]
387            .into_iter()
388            .map(|id| SplitId::new(id).unwrap())
389            .collect();
390        let batch: Vec<(SplitId, SplitProgress)> = ids
391            .iter()
392            .zip(1..)
393            .map(|(id, w)| (id.clone(), SplitProgress::new(w, vec![])))
394            .collect();
395        script.fail_next_commit(&ids[1], CoordinationErrorKind::Retryable);
396
397        let results = coordinator.commit_final(&batch);
398
399        assert_eq!(script.final_commits(), vec![batch.clone()]);
400        assert_eq!(results.len(), 3);
401        assert!(results[0].is_ok());
402        assert_eq!(
403            results[1].as_ref().unwrap_err().kind,
404            CoordinationErrorKind::Retryable
405        );
406        assert!(results[2].is_ok());
407        assert_eq!(script.commits(), vec![batch[0].clone(), batch[2].clone()]);
408    }
409}
spate_test/coordination.rs

spate_test/
coordination.rs
