spate_test/
sink.rs

1//! Sink-side mocks: a scriptable [`ShardWriter`] that captures every write,
2//! and a length-prefixed [`RowEncoder`] whose frames tests can decode.
3
4use bytes::{BufMut, Bytes, BytesMut};
5use spate_core::deser::Owned;
6use spate_core::error::{ErrorClass, SinkError};
7use spate_core::hash::FastMap;
8use spate_core::record::Record;
9use spate_core::sink::{
10    FrameCodecFactory, RowEncoder, SealedBatch, ShardWriter, SinkBundle, SinkMetricsCtx, SinkParts,
11    SinkPoolConfig, endpoint_probe,
12};
13use std::collections::VecDeque;
14use std::sync::{Arc, Mutex, MutexGuard};
15use std::task::{Poll, Waker};
16use std::time::Duration;
17
18/// A replica endpoint tag, [`CaptureWriter`]'s
19/// [`Endpoint`](ShardWriter::Endpoint).
20#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
21pub struct ReplicaTag {
22    /// Shard index.
23    pub shard: usize,
24    /// Replica index within the shard.
25    pub replica: usize,
26}
27
28/// How a scripted write resolves, or that it has not yet.
29#[derive(Clone, Debug, PartialEq, Eq)]
30#[non_exhaustive]
31pub enum ScriptedResult {
32    /// The write succeeds (the durable-ack point).
33    Ok,
34    /// The write fails with [`ErrorClass::Retryable`].
35    Retryable(String),
36    /// The write fails with [`ErrorClass::Fatal`].
37    Fatal(String),
38    /// The write has not resolved: it is held, in its delay, or its future
39    /// was dropped before it resolved. Never a scripted outcome.
40    Pending,
41}
42
43/// One scripted write outcome: an optional hold, an optional delay, then a
44/// result.
45#[derive(Clone, Debug, PartialEq, Eq)]
46#[non_exhaustive]
47pub struct WriteOutcome {
48    /// Wait for [`SinkScript::release_held`] before the delay starts.
49    pub hold: bool,
50    /// Sleep this long (tokio time, controllable with `test-util`) before
51    /// resolving.
52    pub delay: Option<Duration>,
53    /// How the write resolves.
54    pub result: ScriptedResult,
55}
56
57impl WriteOutcome {
58    /// A successful write.
59    #[must_use]
60    pub fn ok() -> Self {
61        WriteOutcome {
62            hold: false,
63            delay: None,
64            result: ScriptedResult::Ok,
65        }
66    }
67
68    /// A retryable failure.
69    #[must_use]
70    pub fn retryable(reason: impl Into<String>) -> Self {
71        WriteOutcome {
72            hold: false,
73            delay: None,
74            result: ScriptedResult::Retryable(reason.into()),
75        }
76    }
77
78    /// A fatal failure.
79    #[must_use]
80    pub fn fatal(reason: impl Into<String>) -> Self {
81        WriteOutcome {
82            hold: false,
83            delay: None,
84            result: ScriptedResult::Fatal(reason.into()),
85        }
86    }
87
88    /// Resolve only after `delay` of tokio time.
89    #[must_use]
90    pub fn after(mut self, delay: Duration) -> Self {
91        self.delay = Some(delay);
92        self
93    }
94
95    /// Keep the write open until the test calls
96    /// [`SinkScript::release_held`].
97    ///
98    /// No amount of time releases a held write, real or tokio, so a test can
99    /// assert what the pipeline does while a write is outstanding without
100    /// reading a clock. A shutdown waits on a held write like any other, so
101    /// release it before stopping the run.
102    #[must_use]
103    pub fn held(mut self) -> Self {
104        self.hold = true;
105        self
106    }
107}
108
109/// One recorded [`ShardWriter::write_batch`] attempt.
110#[derive(Clone, Debug)]
111pub struct CapturedWrite {
112    /// Target shard.
113    pub shard: usize,
114    /// Target replica.
115    pub replica: usize,
116    /// `SealedBatch::rows`.
117    pub rows: u64,
118    /// `SealedBatch::bytes`.
119    pub bytes: u64,
120    /// `SealedBatch::dedup_token`.
121    pub dedup_token: String,
122    /// All frames concatenated (decode with [`decode_rows`](crate::decode_rows)
123    /// when produced by [`TestEncoder`]).
124    pub payload: Vec<u8>,
125    /// How the attempt resolved, or [`Pending`](ScriptedResult::Pending)
126    /// until it does.
127    pub result: ScriptedResult,
128}
129
130#[derive(Debug, Default)]
131struct SinkState {
132    global: VecDeque<WriteOutcome>,
133    per_endpoint: FastMap<(usize, usize), VecDeque<WriteOutcome>>,
134    probe_failures: FastMap<(usize, usize), String>,
135    writes: Vec<CapturedWrite>,
136    /// Count of `release_held` calls. A held write waits for this to pass
137    /// the count it read when it started.
138    releases: u64,
139    /// Writes waiting on a release.
140    holding: usize,
141    /// Wakers of the writes waiting on a release.
142    held_wakers: Vec<Waker>,
143}
144
145/// Counts a write in [`SinkState::holding`] for as long as it waits,
146/// including when the write future is dropped mid-hold.
147struct Holding(Arc<Mutex<SinkState>>);
148
149impl Drop for Holding {
150    fn drop(&mut self) {
151        if let Ok(mut st) = self.0.lock() {
152            st.holding -= 1;
153        }
154    }
155}
156
157/// Refuse an outcome scripted to resolve as `Pending`, at the test's own
158/// call rather than inside a pipeline's sink task.
159fn refuse_pending(outcome: &WriteOutcome) {
160    assert!(
161        outcome.result != ScriptedResult::Pending,
162        "a write cannot be scripted to resolve as Pending; script it with WriteOutcome::held()"
163    );
164}
165
166/// Build a paired scriptable writer and its handle.
167#[must_use]
168pub fn capture_writer() -> (CaptureWriter, SinkScript) {
169    let state = Arc::new(Mutex::new(SinkState::default()));
170    (
171        CaptureWriter {
172            state: Arc::clone(&state),
173        },
174        SinkScript { state },
175    )
176}
177
178/// Scripts outcomes for a [`CaptureWriter`] and reads back its captures.
179/// Cloneable; all clones address the same writer.
180#[derive(Clone, Debug)]
181pub struct SinkScript {
182    state: Arc<Mutex<SinkState>>,
183}
184
185impl SinkScript {
186    fn lock(&self) -> MutexGuard<'_, SinkState> {
187        self.state.lock().expect("spate-test sink state poisoned")
188    }
189
190    /// Queue an outcome consumed by the next otherwise-unscripted write to
191    /// **any** endpoint (FIFO). Endpoint-specific scripts take precedence.
192    ///
193    /// # Panics
194    ///
195    /// If `outcome.result` is [`Pending`](ScriptedResult::Pending).
196    pub fn enqueue_global(&self, outcome: WriteOutcome) {
197        refuse_pending(&outcome);
198        self.lock().global.push_back(outcome);
199    }
200
201    /// Queue an outcome for the next write to `(shard, replica)` (FIFO).
202    ///
203    /// # Panics
204    ///
205    /// If `outcome.result` is [`Pending`](ScriptedResult::Pending).
206    pub fn enqueue_for(&self, shard: usize, replica: usize, outcome: WriteOutcome) {
207        refuse_pending(&outcome);
208        self.lock()
209            .per_endpoint
210            .entry((shard, replica))
211            .or_default()
212            .push_back(outcome);
213    }
214
215    /// Make [`ShardWriter::probe`] of `(shard, replica)` fail with a
216    /// retryable error until [`heal_probe`](Self::heal_probe).
217    pub fn fail_probe(&self, shard: usize, replica: usize, reason: impl Into<String>) {
218        self.lock()
219            .probe_failures
220            .insert((shard, replica), reason.into());
221    }
222
223    /// Let `(shard, replica)`'s probe succeed again.
224    pub fn heal_probe(&self, shard: usize, replica: usize) {
225        self.lock().probe_failures.remove(&(shard, replica));
226    }
227
228    /// Every write attempt so far, in the order they started.
229    ///
230    /// An attempt reads [`Pending`](ScriptedResult::Pending) until its write
231    /// resolves.
232    #[must_use]
233    pub fn writes(&self) -> Vec<CapturedWrite> {
234        self.lock().writes.clone()
235    }
236
237    /// Release every write currently held by a
238    /// [`held`](WriteOutcome::held) outcome.
239    ///
240    /// A held outcome that a later write consumes holds again until the next
241    /// call.
242    pub fn release_held(&self) {
243        let wakers = {
244            let mut st = self.lock();
245            st.releases += 1;
246            std::mem::take(&mut st.held_wakers)
247        };
248        for waker in wakers {
249            waker.wake();
250        }
251    }
252
253    /// How many writes are waiting on [`release_held`](Self::release_held).
254    #[must_use]
255    pub fn held_writes(&self) -> usize {
256        self.lock().holding
257    }
258}
259
260/// A [`ShardWriter`] that records every attempt and resolves it according
261/// to the paired [`SinkScript`] (default: success). Unscripted writes
262/// succeed, so happy-path tests need no scripting.
263#[derive(Clone, Debug)]
264pub struct CaptureWriter {
265    state: Arc<Mutex<SinkState>>,
266}
267
268impl ShardWriter for CaptureWriter {
269    type Endpoint = ReplicaTag;
270
271    fn attach_metrics(&mut self, cx: SinkMetricsCtx) {
272        // Exercise the connector-metrics seam: a family under the sink's
273        // `spate_<component_type>_sink_*` namespace (e.g. `spate_capture_sink_*`).
274        if let Some(meter) = cx.meter {
275            meter.counter("attaches_total", &[]).increment(1);
276        }
277    }
278
279    fn write_batch(
280        &self,
281        endpoint: &ReplicaTag,
282        batch: &SealedBatch,
283    ) -> impl Future<Output = Result<(), SinkError>> + Send {
284        let state = Arc::clone(&self.state);
285        let endpoint = *endpoint;
286        let rows = batch.rows;
287        let bytes = batch.bytes;
288        let dedup_token = batch.dedup_token.clone();
289        let payload: Vec<u8> = batch
290            .frames
291            .iter()
292            .flat_map(|f: &Bytes| f.iter().copied())
293            .collect();
294        async move {
295            let (outcome, hold, at) = {
296                let mut st = state.lock().expect("spate-test sink state poisoned");
297                let outcome = st
298                    .per_endpoint
299                    .get_mut(&(endpoint.shard, endpoint.replica))
300                    .and_then(VecDeque::pop_front)
301                    .or_else(|| st.global.pop_front())
302                    .unwrap_or_else(WriteOutcome::ok);
303                let at = st.writes.len();
304                st.writes.push(CapturedWrite {
305                    shard: endpoint.shard,
306                    replica: endpoint.replica,
307                    rows,
308                    bytes,
309                    dedup_token,
310                    payload,
311                    result: ScriptedResult::Pending,
312                });
313                // Read under the lock that recorded the write, so a release
314                // that follows a test's read of `writes()` reaches this one.
315                let hold = outcome.hold.then(|| {
316                    st.holding += 1;
317                    (st.releases, Holding(Arc::clone(&state)))
318                });
319                (outcome, hold, at)
320            };
321            if let Some((started_at, _holding)) = hold {
322                std::future::poll_fn(|cx| {
323                    let mut st = state.lock().expect("spate-test sink state poisoned");
324                    if st.releases > started_at {
325                        Poll::Ready(())
326                    } else {
327                        st.held_wakers.push(cx.waker().clone());
328                        Poll::Pending
329                    }
330                })
331                .await;
332            }
333            if let Some(delay) = outcome.delay {
334                tokio::time::sleep(delay).await;
335            }
336            let resolved = match &outcome.result {
337                ScriptedResult::Ok => Ok(()),
338                ScriptedResult::Retryable(reason) => Err(SinkError::Client {
339                    class: ErrorClass::Retryable,
340                    reason: reason.clone(),
341                }),
342                ScriptedResult::Fatal(reason) => Err(SinkError::Client {
343                    class: ErrorClass::Fatal,
344                    reason: reason.clone(),
345                }),
346                ScriptedResult::Pending => {
347                    unreachable!("the enqueue methods refuse a Pending outcome")
348                }
349            };
350            // Recorded before the future returns, so a test that reads
351            // `writes()` after anything the return causes sees the result.
352            // `writes` only grows, so `at` still names this attempt.
353            state.lock().expect("spate-test sink state poisoned").writes[at].result =
354                outcome.result;
355            resolved
356        }
357    }
358
359    fn probe(&self, endpoint: &ReplicaTag) -> impl Future<Output = Result<(), SinkError>> + Send {
360        let state = Arc::clone(&self.state);
361        let endpoint = *endpoint;
362        async move {
363            let failure = state
364                .lock()
365                .expect("spate-test sink state poisoned")
366                .probe_failures
367                .get(&(endpoint.shard, endpoint.replica))
368                .cloned();
369            match failure {
370                None => Ok(()),
371                Some(reason) => Err(SinkError::Client {
372                    class: ErrorClass::Retryable,
373                    reason,
374                }),
375            }
376        }
377    }
378}
379
380/// Build a ready-to-run capturing sink: a `shards × replicas` in-memory
381/// topology over a [`CaptureWriter`], with its scripting handle.
382///
383/// The returned [`CaptureSink`] implements
384/// [`SinkBundle`](spate_core::sink::SinkBundle), so it drops straight into
385/// `Pipeline::sink`, the first-class mock path for testing whole
386/// assemblies. Unscripted writes succeed; drive failures and probe health
387/// through the [`SinkScript`].
388#[must_use]
389pub fn capture_sink(shards: usize, replicas: usize) -> (CaptureSink, SinkScript) {
390    assert!(shards > 0, "capture_sink needs at least one shard");
391    assert!(replicas > 0, "capture_sink needs at least one replica");
392    let (writer, script) = capture_writer();
393    (
394        CaptureSink {
395            writer,
396            shards,
397            replicas,
398            pool: SinkPoolConfig::default(),
399            codec: None,
400        },
401        script,
402    )
403}
404
405/// A capturing sink bundle for whole-pipeline tests; see [`capture_sink`].
406#[derive(Clone, Debug)]
407pub struct CaptureSink {
408    writer: CaptureWriter,
409    shards: usize,
410    replicas: usize,
411    pool: SinkPoolConfig,
412    codec: Option<FrameCodecFactory>,
413}
414
415impl CaptureSink {
416    /// Override the pool tuning (tests usually shrink `linger`).
417    #[must_use]
418    pub fn with_pool_config(mut self, pool: SinkPoolConfig) -> Self {
419        self.pool = pool;
420        self
421    }
422
423    /// Attach a frame codec, so the captured payloads are the codec's output
424    /// and [`CapturedWrite::bytes`] counts the bytes before it.
425    #[must_use]
426    pub fn with_frame_codec(mut self, codec: FrameCodecFactory) -> Self {
427        self.codec = Some(codec);
428        self
429    }
430}
431
432impl SinkBundle for CaptureSink {
433    type Writer = CaptureWriter;
434
435    fn into_parts(self) -> SinkParts<CaptureWriter> {
436        let endpoints: Vec<Vec<ReplicaTag>> = (0..self.shards)
437            .map(|shard| {
438                (0..self.replicas)
439                    .map(|replica| ReplicaTag { shard, replica })
440                    .collect()
441            })
442            .collect();
443        // A mock scripts probe health directly (via `SinkScript`), so this
444        // reuses the capture writer rather than an independent client set.
445        // The separate-probe-clients rule matters for real connectors with a
446        // live pool, not an in-memory fake.
447        let probe = endpoint_probe(self.writer.clone(), Arc::new(endpoints.clone()));
448        let parts = SinkParts::new(self.writer, endpoints, self.pool)
449            .with_component_type("capture")
450            .with_probe(probe);
451        match self.codec {
452            Some(codec) => parts.with_frame_codec(codec),
453            None => parts,
454        }
455    }
456}
457
458/// [`RowEncoder`] writing each row as `u32` little-endian length + bytes.
459/// Decode captured frames with [`decode_rows`](crate::decode_rows).
460#[derive(Clone, Copy, Debug, Default)]
461pub struct TestEncoder;
462
463impl RowEncoder<Owned<Vec<u8>>> for TestEncoder {
464    fn encode<'buf>(&mut self, rec: &Record<Vec<u8>>, buf: &mut BytesMut) -> Result<(), SinkError> {
465        let len = u32::try_from(rec.payload.len()).map_err(|_| SinkError::Client {
466            class: ErrorClass::RecordLevel,
467            reason: "row longer than u32::MAX".to_owned(),
468        })?;
469        buf.put_u32_le(len);
470        buf.extend_from_slice(&rec.payload);
471        Ok(())
472    }
473}
474
475/// Decode a frame produced by [`TestEncoder`] back into rows.
476///
477/// # Panics
478///
479/// Panics on truncated input, meaning the frame did not come from
480/// [`TestEncoder`].
481#[must_use]
482pub fn decode_rows(mut bytes: &[u8]) -> Vec<Vec<u8>> {
483    let mut rows = Vec::new();
484    while !bytes.is_empty() {
485        assert!(bytes.len() >= 4, "decode_rows: truncated length prefix");
486        let len = u32::from_le_bytes(bytes[..4].try_into().expect("4 bytes")) as usize;
487        bytes = &bytes[4..];
488        assert!(bytes.len() >= len, "decode_rows: truncated row body");
489        rows.push(bytes[..len].to_vec());
490        bytes = &bytes[len..];
491    }
492    rows
493}
spate_test/sink.rs

spate_test/
sink.rs
