klieo-core 0.41.1

Core traits + runtime for the klieo agent framework.
Documentation
//! Background task that periodically evicts orphaned entries
//! from cluster-0.20's `klieo-leaders` and cluster-0.22's
//! `klieo-tenants` KV buckets.
//!
//! When a replica crashes without firing
//! [`crate::LeaderHandle`]'s Drop or [`crate::OwnershipHandle`]'s
//! Drop, entries linger until the bucket's `max_age` TTL evicts
//! them. Operators running buckets without TTL (e.g. NATS
//! `storage: file`) benefit from an opt-in reaper task that
//! scans + evicts based on resume-buffer liveness.
//!
//! Eviction policy: an entry is orphaned if its corresponding
//! resume buffer reports terminal. See ADR-025.

use crate::bus::KvStore;
use crate::error::BusError;
use crate::resume::ResumeBuffer;
use std::sync::Arc;
use std::time::Duration;
use tokio::task::JoinHandle;

/// Per-server reaper handle. Drop aborts the scan task.
pub struct KvReaperHandle {
    task: Option<JoinHandle<()>>,
}

impl Drop for KvReaperHandle {
    fn drop(&mut self) {
        if let Some(task) = self.task.take() {
            task.abort();
        }
    }
}

/// Spawn a reaper task that scans `buckets` every `interval`
/// and evicts entries whose corresponding resume buffer is
/// terminal.
///
/// Keys are interpreted as `{kind}.{stream_id}` (e.g.
/// `a2a.t-1`, `mcp.tok-1`) — the prefix is stripped to
/// recover the stream id, then
/// [`ResumeBuffer::is_terminal`] drives the orphan decision.
///
/// Returns a [`KvReaperHandle`] whose Drop aborts the
/// background task.
pub fn spawn_kv_reaper(
    kv: Arc<dyn KvStore>,
    resume_buffer: Arc<dyn ResumeBuffer>,
    buckets: Vec<String>,
    interval: Duration,
) -> KvReaperHandle {
    let task = tokio::spawn(async move {
        loop {
            tokio::time::sleep(interval).await;
            for bucket in &buckets {
                if let Err(err) = scan_and_evict(&kv, &resume_buffer, bucket).await {
                    tracing::warn!(
                        target: "klieo.reaper",
                        bucket = %bucket,
                        error = %err,
                        "reaper scan failed; will retry next interval",
                    );
                }
            }
        }
    });
    KvReaperHandle { task: Some(task) }
}

async fn scan_and_evict(
    kv: &Arc<dyn KvStore>,
    buffer: &Arc<dyn ResumeBuffer>,
    bucket: &str,
) -> Result<(), BusError> {
    let keys = kv.keys(bucket).await?;
    let mut evicted = 0usize;
    for key in keys {
        let stream_id = match key.split_once('.') {
            Some((_, rest)) => rest.to_string(),
            None => continue,
        };
        // Treat errored buffer state as terminal (conservative;
        // defer to the bucket's max_age for edge cases). A buffer
        // that does not track terminal state returns Ok(false) per
        // the ResumeBuffer trait default and is therefore preserved.
        let is_terminal = buffer.is_terminal(&stream_id).await.unwrap_or(true);
        if !is_terminal {
            continue;
        }
        tracing::debug!(
            target: "klieo.reaper",
            bucket = %bucket,
            key = %key,
            "evicting orphaned KV entry",
        );
        match kv.delete(bucket, &key).await {
            Ok(()) => evicted += 1,
            Err(err) => tracing::warn!(
                target: "klieo.reaper",
                bucket = %bucket,
                key = %key,
                error = %err,
                "evict failed; will retry next interval",
            ),
        }
    }
    if evicted > 0 {
        tracing::info!(
            target: "klieo.reaper",
            bucket = %bucket,
            evicted = evicted,
            "reaper scan complete",
        );
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::resume::{ResumeBuffer, ResumeError};
    use crate::test_utils::fake_kv;
    use async_trait::async_trait;
    use bytes::Bytes;
    use futures_core::Stream;
    use std::pin::Pin;

    const TEST_BUCKET: &str = "klieo-leaders";
    const SCAN_INTERVAL: Duration = Duration::from_millis(50);
    const STARTUP_DELAY: Duration = Duration::from_millis(20);

    struct AlwaysTerminal;

    #[async_trait]
    impl ResumeBuffer for AlwaysTerminal {
        async fn record(
            &self,
            _stream_id: &str,
            _event_id: u64,
            _payload: Bytes,
        ) -> Result<(), ResumeError> {
            Ok(())
        }
        async fn replay(
            &self,
            _stream_id: &str,
            _since_id: u64,
        ) -> Result<Pin<Box<dyn Stream<Item = (u64, Bytes)> + Send>>, ResumeError> {
            Err(ResumeError::NotFound("test".into()))
        }
        async fn close(&self, _stream_id: &str) -> Result<(), ResumeError> {
            Ok(())
        }
        async fn is_terminal(&self, _stream_id: &str) -> Result<bool, ResumeError> {
            Ok(true)
        }
    }

    struct AlwaysLive;

    #[async_trait]
    impl ResumeBuffer for AlwaysLive {
        async fn record(
            &self,
            _stream_id: &str,
            _event_id: u64,
            _payload: Bytes,
        ) -> Result<(), ResumeError> {
            Ok(())
        }
        async fn replay(
            &self,
            _stream_id: &str,
            _since_id: u64,
        ) -> Result<Pin<Box<dyn Stream<Item = (u64, Bytes)> + Send>>, ResumeError> {
            Err(ResumeError::NotFound("test".into()))
        }
        async fn close(&self, _stream_id: &str) -> Result<(), ResumeError> {
            Ok(())
        }
        async fn is_terminal(&self, _stream_id: &str) -> Result<bool, ResumeError> {
            Ok(false)
        }
    }

    #[tokio::test]
    async fn reaper_handle_drop_aborts_task() {
        let kv = fake_kv();
        let buffer: Arc<dyn ResumeBuffer> = Arc::new(AlwaysLive);
        let handle = spawn_kv_reaper(kv, buffer, vec![TEST_BUCKET.into()], SCAN_INTERVAL);
        tokio::time::sleep(STARTUP_DELAY).await;
        drop(handle);
        // Drop must not panic; aborted task leaks nothing.
    }

    #[tokio::test]
    async fn scan_evicts_when_buffer_terminal() {
        let kv = fake_kv();
        kv.put(TEST_BUCKET, "a2a.t-orphan", Bytes::from_static(b"alice"))
            .await
            .unwrap();
        assert!(kv.get(TEST_BUCKET, "a2a.t-orphan").await.unwrap().is_some());

        let buffer: Arc<dyn ResumeBuffer> = Arc::new(AlwaysTerminal);
        scan_and_evict(&kv, &buffer, TEST_BUCKET).await.unwrap();

        assert!(
            kv.get(TEST_BUCKET, "a2a.t-orphan").await.unwrap().is_none(),
            "reaper must evict entries whose buffer is terminal",
        );
    }

    #[tokio::test]
    async fn scan_keeps_when_buffer_live() {
        let kv = fake_kv();
        kv.put(TEST_BUCKET, "a2a.t-live", Bytes::from_static(b"bob"))
            .await
            .unwrap();

        let buffer: Arc<dyn ResumeBuffer> = Arc::new(AlwaysLive);
        scan_and_evict(&kv, &buffer, TEST_BUCKET).await.unwrap();

        assert!(
            kv.get(TEST_BUCKET, "a2a.t-live").await.unwrap().is_some(),
            "reaper must NOT evict entries whose buffer is non-terminal",
        );
    }
}