use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{mpsc, oneshot};
use tracing::Instrument as _;
use crate::backend::local::LocalBackend;
use crate::config::DurableConfig;
use crate::error::DurableError;
use crate::ids::JournalSeq;
use crate::journal::{Journal, JournalEntry};
const CHANNEL_CAPACITY: usize = 1024;
const MAX_BATCH: usize = 256;
pub(crate) enum JournalMsg {
AppendBuffered(JournalEntry),
AppendAcked(
JournalEntry,
oneshot::Sender<Result<JournalSeq, DurableError>>,
),
Flush(oneshot::Sender<()>),
}
#[derive(Debug)]
pub struct JournalWriter {
backend: Arc<LocalBackend>,
rx: mpsc::Receiver<JournalMsg>,
flush_interval: Duration,
max_batch: usize,
}
impl JournalWriter {
#[must_use]
pub fn new(backend: Arc<LocalBackend>, config: &DurableConfig) -> (Self, JournalWriterHandle) {
let (tx, rx) = mpsc::channel(CHANNEL_CAPACITY);
let handle = JournalWriterHandle {
tx,
ack_timeout: Duration::from_millis(config.journal_ack_timeout_ms),
};
let writer = Self {
backend,
rx,
flush_interval: Duration::from_millis(config.journal_flush_interval_ms.max(1)),
max_batch: MAX_BATCH,
};
(writer, handle)
}
#[tracing::instrument(name = "durable.writer.run", skip_all)]
pub async fn run(mut self) {
let resume = match self.backend.max_seq().await {
Ok(seq) => seq,
Err(error) => {
tracing::error!(%error, "journal writer could not read resume seq; starting at 0");
None
}
};
tracing::info!(
resume_seq = resume.map(JournalSeq::value),
"journal writer started"
);
let mut buffer: Vec<JournalEntry> = Vec::new();
let mut flush = tokio::time::interval(self.flush_interval);
flush.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
loop {
let keep_running = async {
tokio::select! {
maybe_msg = self.rx.recv() => match maybe_msg {
Some(JournalMsg::AppendBuffered(entry)) => {
buffer.push(entry);
if buffer.len() >= self.max_batch {
self.flush_buffer(&mut buffer).await;
}
true
}
Some(JournalMsg::AppendAcked(entry, reply)) => {
self.flush_buffer(&mut buffer).await;
let result = self.backend.append(entry).await;
let _ = reply.send(result);
true
}
Some(JournalMsg::Flush(reply)) => {
self.flush_buffer(&mut buffer).await;
let _ = reply.send(());
true
}
None => {
self.flush_buffer(&mut buffer).await;
false
}
},
_ = flush.tick() => {
self.flush_buffer(&mut buffer).await;
true
}
}
}
.instrument(tracing::info_span!("durable.writer.run.iter"))
.await;
if !keep_running {
break;
}
}
tracing::info!("journal writer stopped");
}
#[tracing::instrument(name = "durable.writer.flush_buffer", skip_all, fields(batch_size = buffer.len()))]
async fn flush_buffer(&self, buffer: &mut Vec<JournalEntry>) {
if buffer.is_empty() {
return;
}
let depth = u32::try_from(buffer.len()).unwrap_or(u32::MAX);
metrics::gauge!("durable.journal.writer.queue_depth").set(f64::from(depth));
if let Err(error) = self.backend.append_batch(buffer).await {
tracing::warn!(
%error,
dropped = buffer.len(),
"journal group-commit failed; buffered entries dropped (re-run safely on resume)"
);
}
buffer.clear();
}
}
#[derive(Clone, Debug)]
pub struct JournalWriterHandle {
tx: mpsc::Sender<JournalMsg>,
ack_timeout: Duration,
}
impl JournalWriterHandle {
pub fn append_buffered(&self, entry: JournalEntry) {
match self.tx.try_send(JournalMsg::AppendBuffered(entry)) {
Ok(()) => {}
Err(mpsc::error::TrySendError::Full(_)) => {
tracing::warn!(
"journal writer channel full; dropping buffered entry (re-runs safely on resume)"
);
}
Err(mpsc::error::TrySendError::Closed(_)) => {
tracing::warn!("journal writer stopped; dropping buffered entry");
}
}
}
#[tracing::instrument(name = "durable.writer.append_acked", skip_all)]
pub async fn append_acked(&self, entry: JournalEntry) -> Result<JournalSeq, DurableError> {
let (reply_tx, reply_rx) = oneshot::channel();
let send_and_wait = async {
self.tx
.send(JournalMsg::AppendAcked(entry, reply_tx))
.await
.map_err(|_| DurableError::JournalUnavailable)?;
match reply_rx.await {
Ok(result) => result,
Err(_) => Err(DurableError::JournalUnavailable),
}
};
tokio::time::timeout(self.ack_timeout, send_and_wait)
.await
.unwrap_or(Err(DurableError::JournalUnavailable))
}
#[tracing::instrument(name = "durable.writer.flush", skip_all)]
pub async fn flush(&self) -> Result<(), DurableError> {
let (reply_tx, reply_rx) = oneshot::channel();
let send_and_wait = async {
self.tx
.send(JournalMsg::Flush(reply_tx))
.await
.map_err(|_| DurableError::JournalUnavailable)?;
reply_rx.await.map_err(|_| DurableError::JournalUnavailable)
};
tokio::time::timeout(self.ack_timeout, send_and_wait)
.await
.unwrap_or(Err(DurableError::JournalUnavailable))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::effect::EffectClass;
use crate::ids::{ExecutionId, ExecutionKind, IdempotencyKey, StepId};
use crate::journal::EntryKind;
use bytes::Bytes;
use std::assert_matches;
fn step_result(exec: ExecutionId, step: u32, payload: &[u8]) -> JournalEntry {
let step_id = StepId::new(step);
JournalEntry {
seq: None,
execution_id: exec,
kind: ExecutionKind::AgentTurn,
step_id,
entry: EntryKind::StepResult {
idempotency_key: IdempotencyKey::derive(exec, step_id, b"tool:read"),
payload: Bytes::copy_from_slice(payload),
effect: EffectClass::Idempotent,
payload_version: 1,
},
created_at_ms: 100,
}
}
#[tokio::test]
async fn append_acked_times_out_when_writer_is_stalled() {
let (tx, _rx) = mpsc::channel(4);
let handle = JournalWriterHandle {
tx,
ack_timeout: Duration::from_millis(50),
};
let result = handle
.append_acked(step_result(ExecutionId::new(), 0, b"x"))
.await;
assert_matches!(result, Err(DurableError::JournalUnavailable));
}
#[tokio::test]
async fn append_acked_errors_when_writer_is_gone() {
let (tx, rx) = mpsc::channel(4);
drop(rx);
let handle = JournalWriterHandle {
tx,
ack_timeout: Duration::from_secs(30),
};
let result = handle
.append_acked(step_result(ExecutionId::new(), 0, b"x"))
.await;
assert_matches!(result, Err(DurableError::JournalUnavailable));
}
#[tokio::test]
async fn append_buffered_drops_on_full_channel_without_blocking() {
let (tx, mut rx) = mpsc::channel(2);
let handle = JournalWriterHandle {
tx,
ack_timeout: Duration::from_millis(50),
};
let exec = ExecutionId::new();
handle.append_buffered(step_result(exec, 0, b"a"));
handle.append_buffered(step_result(exec, 1, b"b"));
handle.append_buffered(step_result(exec, 2, b"c"));
let mut received = 0;
while rx.try_recv().is_ok() {
received += 1;
}
assert_eq!(received, 2, "the over-capacity buffered entry is dropped");
}
#[cfg(feature = "sqlite")]
mod with_backend {
use super::*;
use crate::DurableConfig;
use crate::backend::local::LocalBackend;
use std::sync::Arc;
async fn mem_backend() -> Arc<LocalBackend> {
let backend = LocalBackend::open(":memory:", 1_048_576).await.unwrap();
backend.init().await.unwrap();
Arc::new(backend)
}
fn fast_config() -> DurableConfig {
DurableConfig {
journal_flush_interval_ms: 5,
journal_ack_timeout_ms: 2000,
..DurableConfig::default()
}
}
#[tokio::test]
async fn writer_group_commits_buffered_and_acks_exactly_once() {
let backend = mem_backend().await;
let exec = ExecutionId::new();
backend
.open_execution(exec, ExecutionKind::AgentTurn)
.await
.unwrap();
let (writer, handle) = JournalWriter::new(backend.clone(), &fast_config());
let task = tokio::spawn(writer.run());
handle.append_buffered(step_result(exec, 0, b"a"));
handle.append_buffered(step_result(exec, 1, b"b"));
let seq = handle
.append_acked(step_result(exec, 2, b"c"))
.await
.unwrap();
assert!(seq.value() >= 1);
handle.flush().await.unwrap();
let entries = backend.read_execution(exec).await.unwrap();
assert_eq!(
entries.len(),
3,
"all buffered and acked entries are durable"
);
drop(handle);
task.await.unwrap();
}
#[tokio::test]
async fn writer_resumes_from_max_seq_after_restart() {
let backend = mem_backend().await;
let exec = ExecutionId::new();
backend
.open_execution(exec, ExecutionKind::AgentTurn)
.await
.unwrap();
let config = fast_config();
let (writer1, handle1) = JournalWriter::new(backend.clone(), &config);
let task1 = tokio::spawn(writer1.run());
for step in 0..3 {
handle1
.append_acked(step_result(exec, step, b"x"))
.await
.unwrap();
}
drop(handle1);
task1.await.unwrap();
assert_eq!(backend.max_seq().await.unwrap(), Some(JournalSeq::new(3)));
let (writer2, handle2) = JournalWriter::new(backend.clone(), &config);
let task2 = tokio::spawn(writer2.run());
let seq4 = handle2
.append_acked(step_result(exec, 3, b"y"))
.await
.unwrap();
assert_eq!(
seq4.value(),
4,
"resumed appends continue with neither gap nor duplication"
);
drop(handle2);
task2.await.unwrap();
}
}
}