use crate::bus::KvStore;
use crate::error::BusError;
use crate::resume::ResumeBuffer;
use std::sync::Arc;
use std::time::Duration;
use tokio::task::JoinHandle;
pub struct KvReaperHandle {
task: Option<JoinHandle<()>>,
}
impl Drop for KvReaperHandle {
fn drop(&mut self) {
if let Some(task) = self.task.take() {
task.abort();
}
}
}
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,
};
match buffer.is_terminal(&stream_id).await {
Ok(true) => {}
Ok(false) => continue,
Err(err) => {
tracing::warn!(
target: "klieo.reaper",
bucket = %bucket,
key = %key,
error = %err,
"terminality probe failed; preserving entry (defer to bucket TTL)",
);
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);
}
#[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",
);
}
struct TerminalCheckErrors;
#[async_trait]
impl ResumeBuffer for TerminalCheckErrors {
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> {
Err(ResumeError::Backend("kv blip".into()))
}
}
#[tokio::test]
async fn scan_preserves_when_terminal_check_errors() {
let kv = fake_kv();
kv.put(TEST_BUCKET, "a2a.t-uncertain", Bytes::from_static(b"carol"))
.await
.unwrap();
let buffer: Arc<dyn ResumeBuffer> = Arc::new(TerminalCheckErrors);
scan_and_evict(&kv, &buffer, TEST_BUCKET).await.unwrap();
assert!(
kv.get(TEST_BUCKET, "a2a.t-uncertain")
.await
.unwrap()
.is_some(),
"reaper must preserve entries when the terminality check errors",
);
}
#[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",
);
}
}