sync from monorepo @ 2452e92e
This commit is contained in:
@@ -0,0 +1,256 @@
|
||||
//! 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(()),
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user