spate_core/sink/

mod.rs

//! Sink abstraction: pipeline threads encode, shard workers batch and
//! write.
//!
//! The division of labour (ADR-0006):
//!
//! - **Pipeline threads** route each record to a shard, then run the
//!   sink's [`RowEncoder`] inside the chain's terminal stage, accumulating
//!   encoded rows into small [`EncodedChunk`] frames per shard and
//!   `try_send`ing them into bounded per-shard queues. The send never
//!   blocks; a full queue surfaces as backpressure. Routing has two tiers
//!   over one seam: meta-only [`ShardRouter`] (the default
//!   [`KeyHashRouter`] uses the key hash, else a stable partition hash) or
//!   record-aware [`RecordRouter`] for payload-derived shard affinity.
//! - **Shard workers** (tokio tasks) merge chunks from all pipeline
//!   threads into full-size batches, seal on `max_rows` / `max_bytes` /
//!   `linger`, and dispatch up to `inflight.max_per_shard` concurrent
//!   [`ShardWriter::write_batch`] calls rotating across healthy replicas.
//!   Merging at the worker keeps batches large regardless of the pipeline
//!   thread count.
//!
//! A connector implements [`RowEncoder`] (CPU half) and [`ShardWriter`]
//! (I/O half), and may ship a [`RecordRouter`] when the target's sharding
//! is payload-derived and a [`FrameCodec`] when its wire form is a coded
//! version of the encoded rows; the framework owns everything between them.

mod breaker;
mod bundle;
mod config;
mod failures;
mod pool;
#[cfg(test)]
mod pool_tests;
mod queue;
mod retry;
mod worker;

pub use bundle::{SinkBundle, SinkParts};
pub use config::{
    BatchConfig, BreakerConfig, BreakerConfigError, InflightConfig, RetryConfig, RetryConfigError,
    SinkPoolConfig,
};
pub use failures::SinkFailures;
pub use pool::{DrainReport, SinkPool};
pub use queue::{ChunkSendError, ShardQueues, shard_queues};

/// Boxed sink drain hook: budget in, report out. Produced by sink
/// assemblies (wrapping [`SinkPool::drain`]), consumed once at shutdown by
/// the pipeline runtime.
pub type SinkDrainFn = Box<
    dyn FnOnce(std::time::Duration) -> std::pin::Pin<Box<dyn Future<Output = DrainReport> + Send>>
        + Send,
>;

/// Boxed, repeatable sink connectivity probe (readiness). The runtime
/// probes at startup and then periodically, driving the sinks-connected
/// half of `/readyz`.
pub type SinkProbeFn = Box<
    dyn Fn() -> std::pin::Pin<Box<dyn Future<Output = Result<(), SinkError>> + Send>> + Send + Sync,
>;

/// Build a [`SinkProbeFn`] that probes every replica of every shard in
/// `shard_endpoints` (indexed `[shard][replica]`) via
/// [`ShardWriter::probe`]. This is the readiness loop
/// [`SinkParts::with_probe`](crate::sink::SinkParts::with_probe) expects.
/// Back `writer` with an independent probe client set, never the insert
/// clients (see [`SinkParts::probe`](crate::sink::SinkParts)).
pub fn endpoint_probe<W>(
    writer: W,
    shard_endpoints: std::sync::Arc<Vec<Vec<W::Endpoint>>>,
) -> SinkProbeFn
where
    W: ShardWriter + Clone,
{
    Box::new(move || {
        let writer = writer.clone();
        let shard_endpoints = std::sync::Arc::clone(&shard_endpoints);
        Box::pin(async move {
            for shard in shard_endpoints.iter() {
                for endpoint in shard {
                    writer.probe(endpoint).await?;
                }
            }
            Ok(())
        })
    })
}

use crate::checkpoint::AckSet;
use crate::deser::RecFamily;
use crate::error::SinkError;
use crate::metrics::Meter;
use crate::record::{Record, RecordMeta};
use bytes::{Bytes, BytesMut};
use std::time::Instant;

/// A small frame of encoded rows produced on a pipeline thread, the unit
/// shipped over the per-shard queues. Wire frames are concatenable, so
/// workers accumulate chunks without re-encoding. A frame is one of three
/// things:
///
/// - headerless rows appended back-to-back (RowBinary);
/// - one complete, self-describing block (ClickHouse Native);
/// - a unit the sink's [`FrameCodec`] made from one of those, which decodes
///   on its own.
///
/// In every case, frames concatenated in seal order form a legal body.
///
/// Teardown safety: `acks` is an [`AckSet`]. Dropping a chunk anywhere
/// (a closed queue, an aborted worker, a parked chunk at teardown) fails
/// its batches so their offsets never commit; only a completed durable
/// write delivers them.
#[derive(Debug)]
pub struct EncodedChunk {
    /// Encoded rows in the sink's wire format, after the sink's
    /// [`FrameCodec`] when it has one.
    pub frame: Bytes,
    /// The frame's length as the encoder wrote it, before any codec. Equal
    /// to `frame.len()` when the sink has no codec. The in-flight budget and
    /// [`BatchConfig::max_bytes`] count this length.
    pub raw_len: usize,
    /// Number of rows in `frame`.
    pub rows: u32,
    /// Acknowledgment handles of the source batches represented in
    /// `frame`. Consecutive records usually share a batch, so this stays
    /// short (the encoder dedupes consecutive identical handles).
    pub acks: AckSet,
    /// When the oldest record in `frame` entered the terminal stage
    /// (ingest-basis end-to-end latency).
    pub oldest_ingest: Instant,
    /// Smallest record event time in `frame`, milliseconds since the epoch
    /// (event-basis end-to-end latency).
    pub oldest_event_ms: i64,
}

/// The CPU half of a sink connector: encodes one record into the sink's
/// wire format. Runs on pinned pipeline threads inside the chain's
/// terminal stage; must not perform I/O. Family-generic and dyn-compatible,
/// like [`Deserializer`](crate::deser::Deserializer).
pub trait RowEncoder<F: RecFamily>: Send {
    /// Append `rec`'s encoding to `buf`. Errors are record-level and
    /// subject to the sink stage's `ErrorPolicy`. Errors of
    /// [`ErrorClass::Fatal`](crate::error::ErrorClass::Fatal) stop the
    /// pipeline regardless of policy; fatal means the encoder itself is
    /// broken (e.g. the row type cannot match the target schema), so every
    /// subsequent record would fail identically.
    fn encode<'buf>(
        &mut self,
        rec: &Record<F::Rec<'buf>>,
        buf: &mut BytesMut,
    ) -> Result<(), SinkError>;

    /// Bytes the encoder is holding internally that have **not** yet been
    /// flushed to a frame. Row formats append directly in
    /// [`encode`](Self::encode) and buffer nothing, so the default is `0`.
    /// Columnar formats (which must transpose a whole block before any bytes
    /// exist) return the approximate size of the block under assembly; the
    /// terminal stage adds this to the shard buffer length when deciding
    /// whether to seal a chunk, so a columnar block still respects
    /// [`ChunkConfig::target_bytes`](crate::ops::ChunkConfig).
    fn buffered_bytes(&self) -> usize {
        0
    }

    /// Finalize the pending chunk: flush any internally-buffered rows into
    /// `buf` as exactly **one** complete, self-describing wire frame, leaving
    /// the encoder empty and ready for the next chunk. Row formats already
    /// wrote every row in [`encode`](Self::encode), so the default is a
    /// no-op. The terminal stage calls this immediately before it seals each
    /// [`EncodedChunk`] (in steady state, on data lulls, and at drain), so a
    /// columnar encoder's buffered rows are not silently dropped.
    ///
    /// An `Err` is fatal (a broken encoder, not a bad record). The stage
    /// discards whatever the failed call wrote and ships no partial frame.
    /// The stop it triggers is asynchronous, so the chunk stays open and a
    /// later seal may call `finish_chunk` again for it, by which point the
    /// chunk may hold rows the failed call never saw. Return `Ok` only from a
    /// call whose frame carries every row the chunk holds. The stage discards
    /// what the call wrote and not the row count, so an encoder that had
    /// already moved its rows into `buf` when it failed re-encodes them, and
    /// keeps failing until it can. Rows no call finalizes have their acknowledgments fail on
    /// teardown, so the data replays.
    ///
    /// Because a Native block concatenates with the blocks around it, each
    /// `finish_chunk` frame is independently valid, so workers still
    /// accumulate frames without re-encoding.
    fn finish_chunk(&mut self, buf: &mut BytesMut) -> Result<(), SinkError> {
        let _ = buf;
        Ok(())
    }
}

/// Turns each sealed chunk into its wire form, on the pipeline thread that
/// sealed it. Compression is the usual case.
///
/// The terminal stage calls [`encode_frame`](Self::encode_frame) once per
/// sealed chunk, after the [`RowEncoder`] has finished it and before the
/// chunk is queued. The output becomes the chunk's frame, is stored, and is
/// never recomputed, so every write attempt of a batch, including retries,
/// sends the same bytes. The in-flight budget and
/// [`BatchConfig::max_bytes`] count the chunk's length before the codec, so
/// neither depends on what a codec does.
///
/// The contract:
///
/// - The output is one unit that decodes on its own, and units concatenated
///   in seal order form a legal stream for the sink. Compressed database
///   blocks and the members of a multi-frame zstd or gzip object body both
///   fit.
/// - The output is a deterministic function of the input.
/// - A codec must not perform I/O.
/// - An instance may keep state between calls, such as a compression
///   context or a scratch buffer. The framework makes one instance per sink
///   per chain, which is one per pipeline thread, and calls it from that
///   thread only.
/// - The terminal stage copies the returned slice before the next call.
/// - An `Err` stops the pipeline whatever its [`ErrorClass`]. Nothing ships
///   for the chunk; its rows and acknowledgments stay in the shard buffer,
///   and they fail on teardown if no later call succeeds.
///
/// The `Sync` bound keeps [`SinkHandoff`], the terminal stage that holds the
/// codec, `Sync`.
///
/// [`ErrorClass`]: crate::error::ErrorClass
/// [`SinkHandoff`]: crate::ops::SinkHandoff
///
/// # Examples
///
/// A codec that writes each chunk's length before it:
///
/// ```
/// use spate_core::error::SinkError;
/// use spate_core::sink::{FrameCodec, FrameCodecFactory};
///
/// #[derive(Default)]
/// struct LengthPrefix {
///     out: Vec<u8>,
/// }
///
/// impl FrameCodec for LengthPrefix {
///     fn encode_frame(&mut self, chunk: &[u8]) -> Result<&[u8], SinkError> {
///         self.out.clear();
///         let len = u32::try_from(chunk.len()).map_err(|_| SinkError::Client {
///             class: spate_core::ErrorClass::Fatal,
///             reason: "chunk longer than u32::MAX".into(),
///         })?;
///         self.out.extend_from_slice(&len.to_le_bytes());
///         self.out.extend_from_slice(chunk);
///         Ok(&self.out)
///     }
/// }
///
/// let factory = FrameCodecFactory::new(|| Box::new(LengthPrefix::default()));
/// let mut codec = factory.make();
/// assert_eq!(codec.encode_frame(b"ab").unwrap(), b"\x02\x00\x00\x00ab");
/// ```
pub trait FrameCodec: Send + Sync {
    /// Code one finished chunk, returning the frame to ship.
    ///
    /// # Errors
    ///
    /// Any error stops the pipeline; see the trait's contract.
    fn encode_frame(&mut self, chunk: &[u8]) -> Result<&[u8], SinkError>;
}

/// Makes the [`FrameCodec`] instances for one sink: one per chain. Cloning
/// shares the factory.
///
/// The closure is `RefUnwindSafe`, so the factory, and a sink type that holds
/// one, stays `UnwindSafe` and `RefUnwindSafe`.
///
/// A closure that captures state which is not `RefUnwindSafe` wraps it in
/// [`AssertUnwindSafe`](std::panic::AssertUnwindSafe) and reaches it through
/// deref, as in `let state: &T = &wrapper;`. From edition 2021 a closure
/// that names `wrapper.0` captures only the inner field, which loses the
/// wrapper and fails the bound.
#[derive(Clone)]
pub struct FrameCodecFactory(std::sync::Arc<dyn MakeFrameCodec>);

/// The closure a [`FrameCodecFactory`] wraps.
trait MakeFrameCodec: Fn() -> Box<dyn FrameCodec> + Send + Sync + std::panic::RefUnwindSafe {}

impl<T> MakeFrameCodec for T where
    T: Fn() -> Box<dyn FrameCodec> + Send + Sync + std::panic::RefUnwindSafe
{
}

impl FrameCodecFactory {
    /// Wrap `make`, which builds one codec instance per call.
    pub fn new(
        make: impl Fn() -> Box<dyn FrameCodec> + Send + Sync + std::panic::RefUnwindSafe + 'static,
    ) -> Self {
        FrameCodecFactory(std::sync::Arc::new(make))
    }

    /// Build a fresh codec instance.
    #[must_use]
    pub fn make(&self) -> Box<dyn FrameCodec> {
        (self.0)()
    }
}

impl std::fmt::Debug for FrameCodecFactory {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("FrameCodecFactory")
    }
}

/// A batch sealed by a shard worker, ready to write. Frames concatenate to
/// the full wire payload (a stream of one or more self-describing blocks for
/// block formats like ClickHouse Native).
#[derive(Debug)]
pub struct SealedBatch {
    /// Encoded frames, in order.
    pub frames: Vec<Bytes>,
    /// Total rows across `frames`.
    pub rows: u64,
    /// Total bytes the encoder wrote for `frames`: the sum of each chunk's
    /// [`EncodedChunk::raw_len`], which exceeds the frames' length when the
    /// sink's [`FrameCodec`] compresses.
    pub bytes: u64,
    /// Deterministic-within-a-session batch identity. Retries of the same
    /// sealed batch (including on other replicas) reuse the same token,
    /// so sinks with server-side deduplication windows treat them as
    /// idempotent. Crash replay produces different tokens (documented
    /// at-least-once semantics).
    pub dedup_token: String,
}

/// What the builder hands a sink writer for its own metric families, through
/// [`ShardWriter::attach_metrics`].
#[derive(Debug)]
#[non_exhaustive]
pub struct SinkMetricsCtx {
    /// Scoped `spate_<component_type>_sink_*` and pre-labeled with the
    /// standard `pipeline`/`component`/`component_type`. `None` when the
    /// sink's `component_type` cannot scope a family: the default `"custom"`
    /// (reserved for pipeline-author metrics) or a reserved root opts out
    /// silently, and a malformed value is logged and also yields `None`.
    /// Declare a distinct `component_type` via
    /// [`SinkParts::with_component_type`](crate::sink::SinkParts::with_component_type)
    /// to receive a scope.
    pub meter: Option<Meter>,
}

impl SinkMetricsCtx {
    /// Context carrying `meter`. Connectors outside this crate construct it
    /// here, as the struct is `#[non_exhaustive]`.
    ///
    /// # Examples
    ///
    /// ```
    /// use spate_core::metrics::{Meter, render_local};
    /// use spate_core::sink::SinkMetricsCtx;
    ///
    /// let meter = Meter::with_namespace("probe_sink", "orders", "sink", "probe");
    /// let cx = SinkMetricsCtx::new(Some(meter));
    /// let meter = cx.meter.expect("a scoped Meter");
    /// let text = render_local(|| meter.counter("writes_total", &[]).increment(1));
    /// assert!(text.contains("spate_probe_sink_writes_total{"));
    /// ```
    #[must_use]
    pub fn new(meter: Option<Meter>) -> Self {
        SinkMetricsCtx { meter }
    }
}

/// The I/O half of a sink connector: writes one sealed batch to one
/// replica endpoint. Returning `Ok` is the durable-ack point; only then
/// may the framework resolve the batch's acknowledgments.
pub trait ShardWriter: Send + Sync + 'static {
    /// A connected replica endpoint (e.g. one HTTP client per replica).
    type Endpoint: Send + Sync + 'static;

    /// Receive the [`SinkMetricsCtx`] for the writer's own metric families.
    /// Called once by the builder before the writer is shared across shard
    /// workers; resolve handles here and store them (they are `Arc`-backed, so
    /// `write_batch`'s `&self` can touch them). See [`SinkMetricsCtx::meter`]
    /// for when it carries no `Meter`. The default ignores it.
    fn attach_metrics(&mut self, _cx: SinkMetricsCtx) {}

    /// Write `batch` to `endpoint` durably.
    ///
    /// **Bound it.** The framework guarantees that shutdown terminates, not
    /// that this call does. At the drain deadline the write task is aborted,
    /// which only lands if this future is at an await point. A client with no
    /// request timeout turns every shutdown into a wait for the deadline, and
    /// one that blocks its thread between awaits cannot be aborted at all.
    /// Give the underlying client a request timeout and let the framework's
    /// retry policy handle the failure.
    fn write_batch(
        &self,
        endpoint: &Self::Endpoint,
        batch: &SealedBatch,
    ) -> impl Future<Output = Result<(), SinkError>> + Send;

    /// Connectivity probe for readiness. Defaults to healthy.
    fn probe(
        &self,
        endpoint: &Self::Endpoint,
    ) -> impl Future<Output = Result<(), SinkError>> + Send {
        let _ = endpoint;
        async { Ok(()) }
    }
}

/// Routes records to shards on metadata alone, the **meta-only tier** of
/// sink routing. Pure and cheap; called per record on pipeline threads.
///
/// Every `ShardRouter` is also a [`RecordRouter`] for every record family
/// through a blanket bridge, so meta-only routers plug into the same
/// builder seam unchanged. Implement [`RecordRouter`] directly instead
/// when routing needs the payload.
pub trait ShardRouter: Send + Sync {
    /// The shard index in `0..num_shards` for a record.
    fn route(&self, meta: &RecordMeta, num_shards: usize) -> usize;
}

/// Default router: key hash modulo shards, falling back to the source
/// partition for keyless records (keeps a partition's keyless records
/// together and the distribution stable).
#[derive(Clone, Copy, Debug, Default)]
pub struct KeyHashRouter;

impl ShardRouter for KeyHashRouter {
    #[inline]
    fn route(&self, meta: &RecordMeta, num_shards: usize) -> usize {
        debug_assert!(num_shards > 0);
        let h = meta
            .key_hash
            .unwrap_or_else(|| u64::from(meta.partition.0).wrapping_mul(0x9E37_79B9_7F4A_7C15));
        (h % num_shards as u64) as usize
    }
}

/// Routes records to shards with access to the full record, the
/// **record-aware tier** of sink routing. Pure and cheap; called once per
/// record on pinned pipeline threads, strictly before encoding. It must
/// not perform I/O, block, or allocate per call. Family-generic and
/// dyn-compatible, like [`RowEncoder`].
///
/// The two tiers share one seam:
///
/// - **Meta-only** ([`ShardRouter`]): routes on [`RecordMeta`] alone (key
///   hash, source partition). The default [`KeyHashRouter`] lives here.
///   Every `ShardRouter` is automatically a `RecordRouter` for every
///   family through a blanket bridge, so meta-only routers plug into the
///   same builder seam unchanged.
/// - **Record-aware** (this trait): routes on the payload itself.
///   Required when shard affinity derives from a field of the terminal
///   record type (e.g. matching a sink cluster's own sharding expression),
///   and the only way to route `flat_map` children independently, since
///   children inherit their parent's [`RecordMeta`] and a meta-only router
///   colocates them.
///
/// The router sees the record exactly as the [`RowEncoder`] will (after
/// every transform), so a routing key must survive to the terminal record
/// type. A router may hold state (a weights table, an atomic counter);
/// `&self` plus interior mutability covers stateful strategies.
///
/// A router must also be **total**: return a shard index for every record
/// and never panic. Routing has no per-record error policy; a record
/// either has a well-defined shard or the router picks a deterministic
/// fallback. Unlike an encoder error, which honors the sink
/// stage's Skip/Fail policy, a router panic fails the in-flight batch and
/// stops the pipeline; restart then replays the same record, so a
/// payload-dependent panic is a deterministic crash loop until a code fix
/// ships.
///
/// # Examples
///
/// A record-aware router over an owned family:
///
/// ```
/// use spate_core::deser::Owned;
/// use spate_core::record::Record;
/// use spate_core::sink::RecordRouter;
///
/// struct ByLen;
/// impl RecordRouter<Owned<Vec<u8>>> for ByLen {
///     fn route_record<'buf>(&self, rec: &Record<Vec<u8>>, num_shards: usize) -> usize {
///         rec.payload.len() % num_shards
///     }
/// }
/// ```
///
/// Implement **either** this trait **or** [`ShardRouter`], never both;
/// the bridge makes implementing both a coherence overlap:
///
/// ```compile_fail,E0119
/// use spate_core::deser::Owned;
/// use spate_core::record::{Record, RecordMeta};
/// use spate_core::sink::{RecordRouter, ShardRouter};
///
/// struct Both;
/// impl ShardRouter for Both {
///     fn route(&self, _: &RecordMeta, _: usize) -> usize { 0 }
/// }
/// impl RecordRouter<Owned<Vec<u8>>> for Both {
///     fn route_record<'buf>(&self, _: &Record<Vec<u8>>, _: usize) -> usize { 0 }
/// }
/// ```
pub trait RecordRouter<F: RecFamily>: Send + Sync {
    /// The shard index in `0..num_shards` for `rec`. Must be
    /// `< num_shards`; the terminal stage indexes its shard buffers
    /// directly with the result.
    fn route_record<'buf>(&self, rec: &Record<F::Rec<'buf>>, num_shards: usize) -> usize;
}

/// Bridge: every meta-only [`ShardRouter`] routes any record family by
/// ignoring the payload and delegating to [`ShardRouter::route`] on the
/// record's metadata.
impl<F: RecFamily, R: ShardRouter> RecordRouter<F> for R {
    #[inline]
    fn route_record<'buf>(&self, rec: &Record<F::Rec<'buf>>, num_shards: usize) -> usize {
        self.route(&rec.meta, num_shards)
    }
}

#[cfg(all(test, not(loom)))]
mod tests {
    use super::*;
    use crate::checkpoint::AckRef;
    use crate::deser::Owned;
    use crate::record::PartitionId;

    fn meta(key_hash: Option<u64>, partition: u32) -> RecordMeta {
        RecordMeta {
            partition: PartitionId(partition),
            offset: 0,
            event_time_ms: 0,
            key_hash,
        }
    }

    #[test]
    fn key_hash_router_uses_key_then_partition() {
        let r = KeyHashRouter;
        assert_eq!(r.route(&meta(Some(10), 0), 4), (10 % 4) as usize);
        // Keyless: stable per partition, and different partitions spread.
        let a = r.route(&meta(None, 0), 4);
        let b = r.route(&meta(None, 0), 4);
        assert_eq!(a, b);
        let spread: crate::hash::FastSet<_> = (0..16).map(|p| r.route(&meta(None, p), 4)).collect();
        assert!(spread.len() > 1, "keyless records must not all colocate");
    }

    #[test]
    fn shard_router_bridges_to_record_router_ignoring_payload() {
        let (ack, _rx) = AckRef::test_pair();
        for key_hash in [Some(10), Some(u64::MAX), None] {
            let rec = Record {
                payload: vec![1u8, 2, 3],
                meta: meta(key_hash, 3),
                ack: ack.clone(),
            };
            for n in [1usize, 2, 4, 7] {
                assert_eq!(
                    RecordRouter::<Owned<Vec<u8>>>::route_record(&KeyHashRouter, &rec, n),
                    ShardRouter::route(&KeyHashRouter, &rec.meta, n),
                    "the bridge must delegate to the meta-only route"
                );
            }
        }
    }

    #[test]
    fn record_router_is_dyn_compatible_for_a_concrete_family() {
        let _router: &dyn RecordRouter<Owned<Vec<u8>>> = &KeyHashRouter;
    }
}