kinetic_buffers/

memory.rs

//! In-memory buffering with backpressure for Kinetic.

use kinetic_core::EventBatch;
use metrics::{Label, counter, gauge};
use std::sync::Arc;
use tokio::sync::mpsc;
use tracing::warn;

/// Policy for what to do when a buffer is full.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
#[serde(rename_all = "snake_case")]
#[derive(Default)]
pub enum WhenFull {
    /// Backpressure upstream (stop accepting messages until space is available).
    #[default]
    Block,
    // Note: Kinetic plan specifies `block` is the only strategy for zero data loss.
    // We explicitly do not implement `DropNewest` or `DropOldest`.
}

/// A sender handle to a memory buffer.
#[derive(Debug, Clone)]
pub struct BufferSender {
    tx: mpsc::Sender<EventBatch>,
    when_full: WhenFull,
    labels: Arc<[Label]>,
}

impl BufferSender {
    /// Sends a batch to the buffer.
    ///
    /// If the buffer is full and `when_full` is `Block`, this will await until space is available,
    /// propagating backpressure upstream.
    pub async fn send(&self, batch: EventBatch) -> Result<(), mpsc::error::SendError<EventBatch>> {
        let events = batch.num_rows();
        let bytes = batch.estimated_size();
        match self.when_full {
            WhenFull::Block => match self.tx.reserve().await {
                Ok(permit) => {
                    permit.send(batch);
                    self.emit_received(events, bytes);
                    Ok(())
                }
                Err(_) => Err(mpsc::error::SendError(batch)),
            },
        }
    }

    /// Returns the current number of free slots in the buffer.
    pub fn capacity(&self) -> usize {
        self.tx.capacity()
    }

    /// Returns the maximum capacity of the buffer.
    pub fn max_capacity(&self) -> usize {
        self.tx.max_capacity()
    }

    /// Returns the utilization of the buffer as a fraction between 0.0 and 1.0.
    pub fn utilization(&self) -> f64 {
        let max = self.tx.max_capacity();
        if max == 0 {
            return 1.0;
        }
        let current = max - self.tx.capacity();
        current as f64 / max as f64
    }

    /// Try to send without blocking. Returns an error if full.
    pub fn try_send(&self, batch: EventBatch) -> Result<(), mpsc::error::TrySendError<EventBatch>> {
        let events = batch.num_rows();
        let bytes = batch.estimated_size();
        let res = self.tx.try_send(batch);
        match &res {
            Ok(_) => {
                self.emit_received(events, bytes);
            }
            Err(mpsc::error::TrySendError::Full(_)) => {
                counter!("buffer_full_events_total", self.labels.iter()).increment(1);
                // Emit BufferEventsDropped if we were load shedding, but we only block.
            }
            _ => {}
        }
        res
    }

    fn emit_received(&self, events: usize, bytes: usize) {
        counter!("buffer_received_events_total", self.labels.iter()).increment(events as u64);
        counter!("buffer_received_bytes_total", self.labels.iter()).increment(bytes as u64);
        gauge!("buffer_size_events", self.labels.iter()).increment(events as f64);
        gauge!("buffer_size_bytes", self.labels.iter()).increment(bytes as f64);
    }
}

/// A receiver handle from a memory buffer.
#[derive(Debug)]
pub struct BufferReceiver {
    rx: mpsc::Receiver<EventBatch>,
    labels: Arc<[Label]>,
}

impl BufferReceiver {
    /// Receives the next batch from the buffer.
    pub async fn recv(&mut self) -> Option<EventBatch> {
        let batch = self.rx.recv().await;
        if let Some(ref b) = batch {
            self.emit_sent(b);
        }
        batch
    }

    /// Tries to receive the next batch from the buffer without blocking.
    pub fn try_recv(&mut self) -> Result<EventBatch, mpsc::error::TryRecvError> {
        let res = self.rx.try_recv();
        if let Ok(ref b) = res {
            self.emit_sent(b);
        }
        res
    }

    fn emit_sent(&self, batch: &EventBatch) {
        let events = batch.num_rows();
        let bytes = batch.estimated_size();
        counter!("buffer_sent_events_total", self.labels.iter()).increment(events as u64);
        counter!("buffer_sent_bytes_total", self.labels.iter()).increment(bytes as u64);
        gauge!("buffer_size_events", self.labels.iter()).decrement(events as f64);
        gauge!("buffer_size_bytes", self.labels.iter()).decrement(bytes as f64);
    }
}

/// Creates a new memory buffer with the given capacity (in batches).
pub fn channel(
    capacity: usize,
    when_full: WhenFull,
    component_id: String,
) -> (BufferSender, BufferReceiver) {
    if when_full != WhenFull::Block {
        warn!(
            "Configured when_full policy {:?} is not supported, falling back to block",
            when_full
        );
    }

    let capacity = capacity.max(1);
    let (tx, rx) = mpsc::channel(capacity);
    let labels: Arc<[Label]> = Arc::new([
        Label::new("component_id", component_id),
        Label::new("buffer_type", "memory"),
    ]);

    // Emit BufferCreated metrics
    gauge!("buffer_max_size_events", labels.iter()).set(capacity as f64);

    (
        BufferSender {
            tx,
            when_full: WhenFull::Block,
            labels: labels.clone(),
        },
        BufferReceiver { rx, labels },
    )
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use arrow_array::RecordBatch;
    use arrow_schema::{DataType, Field, Schema};
    use kinetic_core::{ComponentId, EventMetadata};
    use std::sync::Arc;

    fn empty_batch() -> EventBatch {
        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
        let record_batch = RecordBatch::new_empty(schema);
        let meta = Arc::new(EventMetadata::new("pipe1", ComponentId("src1".to_string())));
        EventBatch::new(record_batch, meta).expect("failed to create batch")
    }

    #[tokio::test]
    async fn test_memory_buffer_block() {
        let (tx, mut rx) = channel(1, WhenFull::Block, "test_buffer".to_string());

        // Fill the buffer
        tx.try_send(empty_batch()).unwrap();

        // Next try_send should fail
        assert!(tx.try_send(empty_batch()).is_err());

        // Recv to free space
        assert!(rx.recv().await.is_some());

        // Now we can send again
        assert!(tx.try_send(empty_batch()).is_ok());
    }
}