kinetic_buffers/

disk.rs

//! Persistent disk-backed buffering for Kinetic.

use arrow_ipc::reader::StreamReader;
use arrow_ipc::writer::StreamWriter;
use kinetic_core::EventBatch;
use metrics::{Label, counter, gauge};
use std::collections::VecDeque;
use std::fs::{self, File};
use std::io::{BufReader, BufWriter};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tracing::{debug, error, info};

/// A persistent buffer that stores Arrow batches on disk.
pub struct DiskBuffer {
    data_dir: PathBuf,
    component_id: String,
    pending_files: VecDeque<PathBuf>,
    labels: Arc<[Label]>,
    pub max_size_bytes: Option<u64>,
    pub current_size_bytes: u64,
}

impl DiskBuffer {
    /// Creates a new DiskBuffer or loads existing state from disk.
    pub fn new<P: AsRef<Path>>(
        data_dir: P,
        component_id: String,
        max_size_bytes: Option<u64>,
    ) -> anyhow::Result<Self> {
        let path = data_dir.as_ref().join(&component_id);
        fs::create_dir_all(&path)?;

        let labels: Arc<[Label]> = Arc::new([
            Label::new("component_id", component_id.clone()),
            Label::new("buffer_type", "disk"),
        ]);

        if let Some(max) = max_size_bytes {
            gauge!("buffer_max_size_bytes", labels.iter()).set(max as f64);
        }

        let mut buffer = Self {
            data_dir: path,
            component_id,
            pending_files: VecDeque::new(),
            labels,
            max_size_bytes,
            current_size_bytes: 0,
        };

        buffer.load_existing()?;
        Ok(buffer)
    }

    /// Loads existing batch files from the data directory.
    fn load_existing(&mut self) -> anyhow::Result<()> {
        let mut files = Vec::new();
        let mut total_size = 0;
        for entry in fs::read_dir(&self.data_dir)? {
            let entry = entry?;
            let path = entry.path();
            if path.is_file() && path.extension().is_some_and(|ext| ext == "arrow") {
                total_size += entry.metadata()?.len();
                files.push(path);
            }
        }

        // Sort by file name (which includes timestamp/xid) to preserve order
        files.sort();
        self.pending_files.extend(files);
        self.current_size_bytes = total_size;

        if !self.pending_files.is_empty() {
            info!(
                "Loaded {} existing batches ({} bytes) for component '{}'",
                self.pending_files.len(),
                self.current_size_bytes,
                self.component_id
            );
        }

        Ok(())
    }

    /// Enqueues a batch to disk.
    pub fn enqueue(&mut self, batch: &EventBatch) -> anyhow::Result<()> {
        if let Some(max) = self.max_size_bytes
            && self.current_size_bytes >= max
        {
            anyhow::bail!(
                "Disk buffer is full ({} >= {} bytes)",
                self.current_size_bytes,
                max
            );
        }

        let events = batch.num_rows();
        let bytes = batch.estimated_size();

        // Generate a unique filename using timestamp and xid
        let timestamp = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)?
            .as_nanos();
        let id = xid::new().to_string();

        let filename = format!("{:020}_{}.arrow", timestamp, id);
        let file_path = self.data_dir.join(filename);

        // Serialize metadata to JSON and store in Arrow schema metadata
        let metadata_json = serde_json::to_string(&batch.metadata)?;
        let mut schema_metadata = batch.payload.schema().metadata().clone();
        schema_metadata.insert("kinetic.metadata".to_string(), metadata_json);

        let new_schema = batch
            .payload
            .schema()
            .as_ref()
            .clone()
            .with_metadata(schema_metadata);
        let rb_with_metadata = arrow_array::RecordBatch::try_new(
            std::sync::Arc::new(new_schema),
            batch.payload.columns().to_vec(),
        )?;

        let file = File::create(&file_path)?;
        let mut writer = StreamWriter::try_new(BufWriter::new(file), &rb_with_metadata.schema())?;
        writer.write(&rb_with_metadata)?;
        writer.finish()?;

        let written_size = fs::metadata(&file_path)?.len();
        self.current_size_bytes += written_size;

        self.pending_files.push_back(file_path);

        // Ack immediately once successfully on disk.
        // This shifts durability from end-to-end to source-to-disk.
        if let Some(token) = batch.ack_token.clone() {
            token.ack();
        }

        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);

        debug!("Enqueued batch to disk: component={}", self.component_id);
        Ok(())
    }

    /// Dequeues the next batch from disk.
    pub fn dequeue(&mut self) -> anyhow::Result<Option<EventBatch>> {
        let Some(file_path) = self.pending_files.pop_front() else {
            return Ok(None);
        };

        let file_size = fs::metadata(&file_path)?.len();

        let batch = {
            let file = File::open(&file_path)?;
            let reader = StreamReader::try_new(BufReader::new(file), None)?;

            // We expect exactly one batch per file in this simple implementation
            let mut record_batches = Vec::new();
            for b in reader {
                record_batches.push(b?);
            }

            if record_batches.is_empty() {
                anyhow::bail!("Arrow file was empty: {:?}", file_path);
            }

            let Some(rb) = record_batches.into_iter().next() else {
                anyhow::bail!("Arrow file had no record batches: {:?}", file_path);
            };

            // Extract metadata from schema
            let schema = rb.schema();
            let metadata_json = schema
                .metadata()
                .get("kinetic.metadata")
                .ok_or_else(|| anyhow::anyhow!("Missing metadata in disk-buffered batch"))?;
            let metadata: kinetic_core::EventMetadata = serde_json::from_str(metadata_json)?;

            EventBatch::new(rb, std::sync::Arc::new(metadata))?
        };

        let events = batch.num_rows();
        let bytes = batch.estimated_size();

        // Clean up the file after successful read
        fs::remove_file(&file_path)?;
        self.current_size_bytes = self.current_size_bytes.saturating_sub(file_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);

        debug!("Dequeued batch from disk: component={}", self.component_id);
        Ok(Some(batch))
    }

    /// Returns the number of pending batches on disk.
    pub fn len(&self) -> usize {
        self.pending_files.len()
    }

    pub fn is_empty(&self) -> bool {
        self.pending_files.is_empty()
    }
}

/// Creates a new disk-backed buffer bridge.
pub fn channel(
    data_dir: PathBuf,
    component_id: String,
    max_size_bytes: Option<u64>,
) -> anyhow::Result<(crate::BufferSender, crate::BufferReceiver)> {
    let disk_buffer = DiskBuffer::new(data_dir, component_id.clone(), max_size_bytes)?;

    // Bridging channels
    let (in_tx, mut in_rx) = crate::channel(
        100,
        crate::WhenFull::Block,
        format!("{}_disk_in", component_id),
    );
    let (out_tx, out_rx) = crate::channel(
        100,
        crate::WhenFull::Block,
        format!("{}_disk_out", component_id),
    );

    let notify = Arc::new(tokio::sync::Notify::new());

    // Create shared state
    let shared_buffer = Arc::new(std::sync::Mutex::new(disk_buffer));

    // Ingress task
    let shared_buffer_in = shared_buffer.clone();
    let notify_in = notify.clone();
    let component_id_in = component_id.clone();
    tokio::spawn(async move {
        while let Some(batch) = in_rx.recv().await {
            // Apply backpressure if full
            loop {
                let is_full = {
                    let b = shared_buffer_in.lock().unwrap_or_else(|p| p.into_inner());
                    match b.max_size_bytes {
                        Some(max) => b.current_size_bytes >= max,
                        None => false,
                    }
                };
                if is_full {
                    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
                    continue;
                }
                break;
            }

            let mut b = shared_buffer_in.lock().unwrap_or_else(|p| p.into_inner());
            if let Err(e) = b.enqueue(&batch) {
                error!(message = "Failed to enqueue batch to disk", %e, component_id = component_id_in);
            } else {
                notify_in.notify_one();
            }
        }
    });

    // Egress task
    let shared_buffer_out = shared_buffer;
    let notify_out = notify;
    let component_id_out = component_id;
    tokio::spawn(async move {
        loop {
            let batch_opt = {
                let mut b = shared_buffer_out.lock().unwrap_or_else(|p| p.into_inner());
                if b.is_empty() {
                    None
                } else {
                    match b.dequeue() {
                        Ok(Some(batch)) => Some(batch),
                        Ok(None) => None,
                        Err(e) => {
                            error!(message = "Failed to dequeue batch from disk", %e, component_id = component_id_out);
                            None
                        }
                    }
                }
            };

            if let Some(batch) = batch_opt {
                if let Err(e) = out_tx.send(batch).await {
                    error!(message = "Failed to send batch from disk to downstream", %e, component_id = component_id_out);
                    return;
                }
            } else {
                notify_out.notified().await;
            }
        }
    });

    Ok((in_tx, out_rx))
}