257 lines
8.9 KiB
Rust
257 lines
8.9 KiB
Rust
//! Per-backend writer task for `WritePolicy::Queued` backends.
|
|
//!
|
|
//! The task drains a per-backend mpsc, optionally batching/coalescing within
|
|
//! a configured window, and invokes `ArchiveBackend` methods directly. Errors
|
|
//! drift health on the parent registration; they do not propagate to the
|
|
//! caller.
|
|
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use chrono::Utc;
|
|
use tokio::sync::{mpsc, oneshot, watch, RwLock};
|
|
use tokio::task::JoinHandle;
|
|
use tracing::{debug, warn};
|
|
use uuid::Uuid;
|
|
|
|
use crate::backend::{ArchiveBackend, HealthStatus};
|
|
|
|
use super::OverflowPolicy;
|
|
|
|
#[derive(Debug)]
|
|
pub enum WriteOp {
|
|
PutSession(crate::types::SessionMetadata),
|
|
AppendMessages {
|
|
scroll_id: Uuid,
|
|
msgs: Vec<crate::types::MessageRecord>,
|
|
},
|
|
DeleteSession {
|
|
scroll_id: Uuid,
|
|
},
|
|
ClearSessionMessages {
|
|
scroll_id: Uuid,
|
|
},
|
|
AppendDagEdge(crate::types::DagEdge),
|
|
AppendMetaEvents {
|
|
scroll_id: Uuid,
|
|
events: Vec<crate::types::MetaEventRecord>,
|
|
},
|
|
Shutdown(oneshot::Sender<()>),
|
|
}
|
|
|
|
impl WriteOp {
|
|
pub fn op_label(&self) -> &'static str {
|
|
match self {
|
|
WriteOp::PutSession(_) => "put_session",
|
|
WriteOp::AppendMessages { .. } => "append_messages",
|
|
WriteOp::DeleteSession { .. } => "delete_session",
|
|
WriteOp::ClearSessionMessages { .. } => "clear_session_messages",
|
|
WriteOp::AppendDagEdge(_) => "append_dag_edge",
|
|
WriteOp::AppendMetaEvents { .. } => "append_meta_events",
|
|
WriteOp::Shutdown(_) => "shutdown",
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct WriterHandle {
|
|
pub sender: mpsc::Sender<WriteOp>,
|
|
pub overflow: OverflowPolicy,
|
|
pub queue_depth: watch::Receiver<usize>,
|
|
pub join: tokio::sync::Mutex<Option<JoinHandle<()>>>,
|
|
pub backend_name: String,
|
|
}
|
|
|
|
impl WriterHandle {
|
|
pub async fn enqueue(&self, op: WriteOp) -> Result<(), crate::error::ArchivistError> {
|
|
match self.overflow {
|
|
OverflowPolicy::Block => self.sender.send(op).await.map_err(|_| {
|
|
crate::error::ArchivistError::Other(format!(
|
|
"writer task for `{}` has closed",
|
|
self.backend_name
|
|
))
|
|
}),
|
|
OverflowPolicy::Error => self.sender.try_send(op).map_err(|e| match e {
|
|
mpsc::error::TrySendError::Full(op) => {
|
|
crate::error::ArchivistError::WriteQueueFull {
|
|
backend: self.backend_name.clone(),
|
|
op: op.op_label(),
|
|
}
|
|
}
|
|
mpsc::error::TrySendError::Closed(_) => {
|
|
crate::error::ArchivistError::Other(format!(
|
|
"writer task for `{}` has closed",
|
|
self.backend_name
|
|
))
|
|
}
|
|
}),
|
|
OverflowPolicy::DropOldest => {
|
|
// Tokio mpsc can't truly "drop oldest" without draining from the
|
|
// other side; we approximate with "drop newest when full". For
|
|
// observability sinks this is acceptable — the contract is
|
|
// "never block, may lose data".
|
|
let _ = self.sender.try_send(op);
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn queue_depth_now(&self) -> usize {
|
|
*self.queue_depth.borrow()
|
|
}
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn spawn_writer(
|
|
backend: Arc<dyn ArchiveBackend>,
|
|
backend_name: String,
|
|
capacity: usize,
|
|
batch_window: Duration,
|
|
overflow: OverflowPolicy,
|
|
health: Arc<RwLock<HealthStatus>>,
|
|
last_error: Arc<RwLock<Option<(chrono::DateTime<chrono::Utc>, String)>>>,
|
|
consecutive_failures: Arc<RwLock<u32>>,
|
|
) -> WriterHandle {
|
|
let (tx, mut rx) = mpsc::channel::<WriteOp>(capacity);
|
|
let (depth_tx, depth_rx) = watch::channel(0usize);
|
|
|
|
let join = tokio::spawn({
|
|
let backend_name = backend_name.clone();
|
|
async move {
|
|
const FAILURE_THRESHOLD: u32 = 5;
|
|
|
|
loop {
|
|
let Some(first) = rx.recv().await else { break };
|
|
let mut batch: Vec<WriteOp> = vec![first];
|
|
|
|
let deadline = tokio::time::Instant::now() + batch_window;
|
|
while tokio::time::Instant::now() < deadline {
|
|
match tokio::time::timeout_at(deadline, rx.recv()).await {
|
|
Ok(Some(op)) => batch.push(op),
|
|
Ok(None) => break,
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
|
|
let _ = depth_tx.send(rx.len());
|
|
|
|
let coalesced = coalesce(batch);
|
|
|
|
let mut shutdown_ack: Option<oneshot::Sender<()>> = None;
|
|
for op in coalesced {
|
|
if let WriteOp::Shutdown(ack) = op {
|
|
shutdown_ack = Some(ack);
|
|
break;
|
|
}
|
|
match dispatch_op(&*backend, op).await {
|
|
Ok(()) => {
|
|
*consecutive_failures.write().await = 0;
|
|
let mut h = health.write().await;
|
|
if matches!(*h, HealthStatus::Degraded { .. }) {
|
|
*h = HealthStatus::Healthy;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
warn!(
|
|
backend = backend_name.as_str(),
|
|
error = %e,
|
|
"queued write failed; drifting health"
|
|
);
|
|
let mut n = consecutive_failures.write().await;
|
|
*n = n.saturating_add(1);
|
|
*last_error.write().await = Some((Utc::now(), format!("{e}")));
|
|
let mut h = health.write().await;
|
|
if *n >= FAILURE_THRESHOLD {
|
|
*h = HealthStatus::Unavailable {
|
|
reason: format!("{} consecutive failures", *n),
|
|
};
|
|
} else {
|
|
*h = HealthStatus::Degraded { reason: format!("{e}") };
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if let Some(ack) = shutdown_ack {
|
|
debug!(backend = backend_name.as_str(), "writer task shutting down");
|
|
let _ = ack.send(());
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
WriterHandle {
|
|
sender: tx,
|
|
overflow,
|
|
queue_depth: depth_rx,
|
|
join: tokio::sync::Mutex::new(Some(join)),
|
|
backend_name,
|
|
}
|
|
}
|
|
|
|
fn coalesce(batch: Vec<WriteOp>) -> Vec<WriteOp> {
|
|
let mut out: Vec<WriteOp> = Vec::with_capacity(batch.len());
|
|
for op in batch {
|
|
let merged = match (out.last_mut(), &op) {
|
|
(
|
|
Some(WriteOp::AppendMessages { scroll_id: a, .. }),
|
|
WriteOp::AppendMessages { scroll_id: b, .. },
|
|
) if a == b => true,
|
|
(
|
|
Some(WriteOp::AppendMetaEvents { scroll_id: a, .. }),
|
|
WriteOp::AppendMetaEvents { scroll_id: b, .. },
|
|
) if a == b => true,
|
|
_ => false,
|
|
};
|
|
|
|
if merged {
|
|
match out.last_mut().unwrap() {
|
|
WriteOp::AppendMessages { msgs: m1, .. } => {
|
|
if let WriteOp::AppendMessages { msgs: m2, .. } = op {
|
|
m1.extend(m2);
|
|
continue;
|
|
}
|
|
}
|
|
WriteOp::AppendMetaEvents { events: e1, .. } => {
|
|
if let WriteOp::AppendMetaEvents { events: e2, .. } = op {
|
|
e1.extend(e2);
|
|
continue;
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
out.push(op);
|
|
}
|
|
out
|
|
}
|
|
|
|
async fn dispatch_op(backend: &dyn ArchiveBackend, op: WriteOp) -> crate::error::Result<()> {
|
|
match op {
|
|
WriteOp::PutSession(meta) => backend.put_session(meta).await,
|
|
WriteOp::AppendMessages { scroll_id, msgs } => {
|
|
backend.append_messages(scroll_id, msgs).await
|
|
}
|
|
WriteOp::DeleteSession { scroll_id } => backend.delete_session(scroll_id).await,
|
|
WriteOp::ClearSessionMessages { scroll_id } => {
|
|
backend.clear_session_messages(scroll_id).await
|
|
}
|
|
WriteOp::AppendDagEdge(edge) => {
|
|
if let Some(d) = backend.as_dag() {
|
|
d.append_dag_edge(edge).await
|
|
} else {
|
|
Ok(())
|
|
}
|
|
}
|
|
WriteOp::AppendMetaEvents { scroll_id, events } => {
|
|
if let Some(m) = backend.as_meta_events() {
|
|
m.append_meta_events(scroll_id, events).await
|
|
} else {
|
|
Ok(())
|
|
}
|
|
}
|
|
WriteOp::Shutdown(_) => Ok(()),
|
|
}
|
|
}
|