spate-test 0.4.0

//! Sink-side mocks: a scriptable [`ShardWriter`] that captures every write,
//! and a length-prefixed [`RowEncoder`] whose frames tests can decode.

use bytes::{BufMut, Bytes, BytesMut};
use spate_core::deser::Owned;
use spate_core::error::{ErrorClass, SinkError};
use spate_core::hash::FastMap;
use spate_core::record::Record;
use spate_core::sink::{
    FrameCodecFactory, RowEncoder, SealedBatch, ShardWriter, SinkBundle, SinkMetricsCtx, SinkParts,
    SinkPoolConfig, endpoint_probe,
};
use std::collections::VecDeque;
use std::sync::{Arc, Mutex, MutexGuard};
use std::task::{Poll, Waker};
use std::time::Duration;

/// A replica endpoint tag, [`CaptureWriter`]'s
/// [`Endpoint`](ShardWriter::Endpoint).
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct ReplicaTag {
    /// Shard index.
    pub shard: usize,
    /// Replica index within the shard.
    pub replica: usize,
}

/// How a scripted write resolves, or that it has not yet.
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum ScriptedResult {
    /// The write succeeds (the durable-ack point).
    Ok,
    /// The write fails with [`ErrorClass::Retryable`].
    Retryable(String),
    /// The write fails with [`ErrorClass::Fatal`].
    Fatal(String),
    /// The write has not resolved: it is held, in its delay, or its future
    /// was dropped before it resolved. Never a scripted outcome.
    Pending,
}

/// One scripted write outcome: an optional hold, an optional delay, then a
/// result.
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct WriteOutcome {
    /// Wait for [`SinkScript::release_held`] before the delay starts.
    pub hold: bool,
    /// Sleep this long (tokio time, controllable with `test-util`) before
    /// resolving.
    pub delay: Option<Duration>,
    /// How the write resolves.
    pub result: ScriptedResult,
}

impl WriteOutcome {
    /// A successful write.
    #[must_use]
    pub fn ok() -> Self {
        WriteOutcome {
            hold: false,
            delay: None,
            result: ScriptedResult::Ok,
        }
    }

    /// A retryable failure.
    #[must_use]
    pub fn retryable(reason: impl Into<String>) -> Self {
        WriteOutcome {
            hold: false,
            delay: None,
            result: ScriptedResult::Retryable(reason.into()),
        }
    }

    /// A fatal failure.
    #[must_use]
    pub fn fatal(reason: impl Into<String>) -> Self {
        WriteOutcome {
            hold: false,
            delay: None,
            result: ScriptedResult::Fatal(reason.into()),
        }
    }

    /// Resolve only after `delay` of tokio time.
    #[must_use]
    pub fn after(mut self, delay: Duration) -> Self {
        self.delay = Some(delay);
        self
    }

    /// Keep the write open until the test calls
    /// [`SinkScript::release_held`].
    ///
    /// No amount of time releases a held write, real or tokio, so a test can
    /// assert what the pipeline does while a write is outstanding without
    /// reading a clock. A shutdown waits on a held write like any other, so
    /// release it before stopping the run.
    #[must_use]
    pub fn held(mut self) -> Self {
        self.hold = true;
        self
    }
}

/// One recorded [`ShardWriter::write_batch`] attempt.
#[derive(Clone, Debug)]
pub struct CapturedWrite {
    /// Target shard.
    pub shard: usize,
    /// Target replica.
    pub replica: usize,
    /// `SealedBatch::rows`.
    pub rows: u64,
    /// `SealedBatch::bytes`.
    pub bytes: u64,
    /// `SealedBatch::dedup_token`.
    pub dedup_token: String,
    /// All frames concatenated (decode with [`decode_rows`](crate::decode_rows)
    /// when produced by [`TestEncoder`]).
    pub payload: Vec<u8>,
    /// How the attempt resolved, or [`Pending`](ScriptedResult::Pending)
    /// until it does.
    pub result: ScriptedResult,
}

#[derive(Debug, Default)]
struct SinkState {
    global: VecDeque<WriteOutcome>,
    per_endpoint: FastMap<(usize, usize), VecDeque<WriteOutcome>>,
    probe_failures: FastMap<(usize, usize), String>,
    writes: Vec<CapturedWrite>,
    /// Count of `release_held` calls. A held write waits for this to pass
    /// the count it read when it started.
    releases: u64,
    /// Writes waiting on a release.
    holding: usize,
    /// Wakers of the writes waiting on a release.
    held_wakers: Vec<Waker>,
}

/// Counts a write in [`SinkState::holding`] for as long as it waits,
/// including when the write future is dropped mid-hold.
struct Holding(Arc<Mutex<SinkState>>);

impl Drop for Holding {
    fn drop(&mut self) {
        if let Ok(mut st) = self.0.lock() {
            st.holding -= 1;
        }
    }
}

/// Refuse an outcome scripted to resolve as `Pending`, at the test's own
/// call rather than inside a pipeline's sink task.
fn refuse_pending(outcome: &WriteOutcome) {
    assert!(
        outcome.result != ScriptedResult::Pending,
        "a write cannot be scripted to resolve as Pending; script it with WriteOutcome::held()"
    );
}

/// Build a paired scriptable writer and its handle.
#[must_use]
pub fn capture_writer() -> (CaptureWriter, SinkScript) {
    let state = Arc::new(Mutex::new(SinkState::default()));
    (
        CaptureWriter {
            state: Arc::clone(&state),
        },
        SinkScript { state },
    )
}

/// Scripts outcomes for a [`CaptureWriter`] and reads back its captures.
/// Cloneable; all clones address the same writer.
#[derive(Clone, Debug)]
pub struct SinkScript {
    state: Arc<Mutex<SinkState>>,
}

impl SinkScript {
    fn lock(&self) -> MutexGuard<'_, SinkState> {
        self.state.lock().expect("spate-test sink state poisoned")
    }

    /// Queue an outcome consumed by the next otherwise-unscripted write to
    /// **any** endpoint (FIFO). Endpoint-specific scripts take precedence.
    ///
    /// # Panics
    ///
    /// If `outcome.result` is [`Pending`](ScriptedResult::Pending).
    pub fn enqueue_global(&self, outcome: WriteOutcome) {
        refuse_pending(&outcome);
        self.lock().global.push_back(outcome);
    }

    /// Queue an outcome for the next write to `(shard, replica)` (FIFO).
    ///
    /// # Panics
    ///
    /// If `outcome.result` is [`Pending`](ScriptedResult::Pending).
    pub fn enqueue_for(&self, shard: usize, replica: usize, outcome: WriteOutcome) {
        refuse_pending(&outcome);
        self.lock()
            .per_endpoint
            .entry((shard, replica))
            .or_default()
            .push_back(outcome);
    }

    /// Make [`ShardWriter::probe`] of `(shard, replica)` fail with a
    /// retryable error until [`heal_probe`](Self::heal_probe).
    pub fn fail_probe(&self, shard: usize, replica: usize, reason: impl Into<String>) {
        self.lock()
            .probe_failures
            .insert((shard, replica), reason.into());
    }

    /// Let `(shard, replica)`'s probe succeed again.
    pub fn heal_probe(&self, shard: usize, replica: usize) {
        self.lock().probe_failures.remove(&(shard, replica));
    }

    /// Every write attempt so far, in the order they started.
    ///
    /// An attempt reads [`Pending`](ScriptedResult::Pending) until its write
    /// resolves.
    #[must_use]
    pub fn writes(&self) -> Vec<CapturedWrite> {
        self.lock().writes.clone()
    }

    /// Release every write currently held by a
    /// [`held`](WriteOutcome::held) outcome.
    ///
    /// A held outcome that a later write consumes holds again until the next
    /// call.
    pub fn release_held(&self) {
        let wakers = {
            let mut st = self.lock();
            st.releases += 1;
            std::mem::take(&mut st.held_wakers)
        };
        for waker in wakers {
            waker.wake();
        }
    }

    /// How many writes are waiting on [`release_held`](Self::release_held).
    #[must_use]
    pub fn held_writes(&self) -> usize {
        self.lock().holding
    }
}

/// A [`ShardWriter`] that records every attempt and resolves it according
/// to the paired [`SinkScript`] (default: success). Unscripted writes
/// succeed, so happy-path tests need no scripting.
#[derive(Clone, Debug)]
pub struct CaptureWriter {
    state: Arc<Mutex<SinkState>>,
}

impl ShardWriter for CaptureWriter {
    type Endpoint = ReplicaTag;

    fn attach_metrics(&mut self, cx: SinkMetricsCtx) {
        // Exercise the connector-metrics seam: a family under the sink's
        // `spate_<component_type>_sink_*` namespace (e.g. `spate_capture_sink_*`).
        if let Some(meter) = cx.meter {
            meter.counter("attaches_total", &[]).increment(1);
        }
    }

    fn write_batch(
        &self,
        endpoint: &ReplicaTag,
        batch: &SealedBatch,
    ) -> impl Future<Output = Result<(), SinkError>> + Send {
        let state = Arc::clone(&self.state);
        let endpoint = *endpoint;
        let rows = batch.rows;
        let bytes = batch.bytes;
        let dedup_token = batch.dedup_token.clone();
        let payload: Vec<u8> = batch
            .frames
            .iter()
            .flat_map(|f: &Bytes| f.iter().copied())
            .collect();
        async move {
            let (outcome, hold, at) = {
                let mut st = state.lock().expect("spate-test sink state poisoned");
                let outcome = st
                    .per_endpoint
                    .get_mut(&(endpoint.shard, endpoint.replica))
                    .and_then(VecDeque::pop_front)
                    .or_else(|| st.global.pop_front())
                    .unwrap_or_else(WriteOutcome::ok);
                let at = st.writes.len();
                st.writes.push(CapturedWrite {
                    shard: endpoint.shard,
                    replica: endpoint.replica,
                    rows,
                    bytes,
                    dedup_token,
                    payload,
                    result: ScriptedResult::Pending,
                });
                // Read under the lock that recorded the write, so a release
                // that follows a test's read of `writes()` reaches this one.
                let hold = outcome.hold.then(|| {
                    st.holding += 1;
                    (st.releases, Holding(Arc::clone(&state)))
                });
                (outcome, hold, at)
            };
            if let Some((started_at, _holding)) = hold {
                std::future::poll_fn(|cx| {
                    let mut st = state.lock().expect("spate-test sink state poisoned");
                    if st.releases > started_at {
                        Poll::Ready(())
                    } else {
                        st.held_wakers.push(cx.waker().clone());
                        Poll::Pending
                    }
                })
                .await;
            }
            if let Some(delay) = outcome.delay {
                tokio::time::sleep(delay).await;
            }
            let resolved = match &outcome.result {
                ScriptedResult::Ok => Ok(()),
                ScriptedResult::Retryable(reason) => Err(SinkError::Client {
                    class: ErrorClass::Retryable,
                    reason: reason.clone(),
                }),
                ScriptedResult::Fatal(reason) => Err(SinkError::Client {
                    class: ErrorClass::Fatal,
                    reason: reason.clone(),
                }),
                ScriptedResult::Pending => {
                    unreachable!("the enqueue methods refuse a Pending outcome")
                }
            };
            // Recorded before the future returns, so a test that reads
            // `writes()` after anything the return causes sees the result.
            // `writes` only grows, so `at` still names this attempt.
            state.lock().expect("spate-test sink state poisoned").writes[at].result =
                outcome.result;
            resolved
        }
    }

    fn probe(&self, endpoint: &ReplicaTag) -> impl Future<Output = Result<(), SinkError>> + Send {
        let state = Arc::clone(&self.state);
        let endpoint = *endpoint;
        async move {
            let failure = state
                .lock()
                .expect("spate-test sink state poisoned")
                .probe_failures
                .get(&(endpoint.shard, endpoint.replica))
                .cloned();
            match failure {
                None => Ok(()),
                Some(reason) => Err(SinkError::Client {
                    class: ErrorClass::Retryable,
                    reason,
                }),
            }
        }
    }
}

/// Build a ready-to-run capturing sink: a `shards × replicas` in-memory
/// topology over a [`CaptureWriter`], with its scripting handle.
///
/// The returned [`CaptureSink`] implements
/// [`SinkBundle`](spate_core::sink::SinkBundle), so it drops straight into
/// `Pipeline::sink`, the first-class mock path for testing whole
/// assemblies. Unscripted writes succeed; drive failures and probe health
/// through the [`SinkScript`].
#[must_use]
pub fn capture_sink(shards: usize, replicas: usize) -> (CaptureSink, SinkScript) {
    assert!(shards > 0, "capture_sink needs at least one shard");
    assert!(replicas > 0, "capture_sink needs at least one replica");
    let (writer, script) = capture_writer();
    (
        CaptureSink {
            writer,
            shards,
            replicas,
            pool: SinkPoolConfig::default(),
            codec: None,
        },
        script,
    )
}

/// A capturing sink bundle for whole-pipeline tests; see [`capture_sink`].
#[derive(Clone, Debug)]
pub struct CaptureSink {
    writer: CaptureWriter,
    shards: usize,
    replicas: usize,
    pool: SinkPoolConfig,
    codec: Option<FrameCodecFactory>,
}

impl CaptureSink {
    /// Override the pool tuning (tests usually shrink `linger`).
    #[must_use]
    pub fn with_pool_config(mut self, pool: SinkPoolConfig) -> Self {
        self.pool = pool;
        self
    }

    /// Attach a frame codec, so the captured payloads are the codec's output
    /// and [`CapturedWrite::bytes`] counts the bytes before it.
    #[must_use]
    pub fn with_frame_codec(mut self, codec: FrameCodecFactory) -> Self {
        self.codec = Some(codec);
        self
    }
}

impl SinkBundle for CaptureSink {
    type Writer = CaptureWriter;

    fn into_parts(self) -> SinkParts<CaptureWriter> {
        let endpoints: Vec<Vec<ReplicaTag>> = (0..self.shards)
            .map(|shard| {
                (0..self.replicas)
                    .map(|replica| ReplicaTag { shard, replica })
                    .collect()
            })
            .collect();
        // A mock scripts probe health directly (via `SinkScript`), so this
        // reuses the capture writer rather than an independent client set.
        // The separate-probe-clients rule matters for real connectors with a
        // live pool, not an in-memory fake.
        let probe = endpoint_probe(self.writer.clone(), Arc::new(endpoints.clone()));
        let parts = SinkParts::new(self.writer, endpoints, self.pool)
            .with_component_type("capture")
            .with_probe(probe);
        match self.codec {
            Some(codec) => parts.with_frame_codec(codec),
            None => parts,
        }
    }
}

/// [`RowEncoder`] writing each row as `u32` little-endian length + bytes.
/// Decode captured frames with [`decode_rows`](crate::decode_rows).
#[derive(Clone, Copy, Debug, Default)]
pub struct TestEncoder;

impl RowEncoder<Owned<Vec<u8>>> for TestEncoder {
    fn encode<'buf>(&mut self, rec: &Record<Vec<u8>>, buf: &mut BytesMut) -> Result<(), SinkError> {
        let len = u32::try_from(rec.payload.len()).map_err(|_| SinkError::Client {
            class: ErrorClass::RecordLevel,
            reason: "row longer than u32::MAX".to_owned(),
        })?;
        buf.put_u32_le(len);
        buf.extend_from_slice(&rec.payload);
        Ok(())
    }
}

/// Decode a frame produced by [`TestEncoder`] back into rows.
///
/// # Panics
///
/// Panics on truncated input, meaning the frame did not come from
/// [`TestEncoder`].
#[must_use]
pub fn decode_rows(mut bytes: &[u8]) -> Vec<Vec<u8>> {
    let mut rows = Vec::new();
    while !bytes.is_empty() {
        assert!(bytes.len() >= 4, "decode_rows: truncated length prefix");
        let len = u32::from_le_bytes(bytes[..4].try_into().expect("4 bytes")) as usize;
        bytes = &bytes[4..];
        assert!(bytes.len() >= len, "decode_rows: truncated row body");
        rows.push(bytes[..len].to_vec());
        bytes = &bytes[len..];
    }
    rows
}