71 lines
2.5 KiB
Rust
71 lines
2.5 KiB
Rust
|
//! Bounded, asynchronous persistence for data-plane invocation telemetry.
|
||
|
|
|
||
|
|
use std::sync::Arc;
|
||
|
|
|
||
|
|
use tokio::sync::{Mutex, mpsc, oneshot};
|
||
|
|
|
||
|
|
use crate::persistence::{InvocationUpdate, Persistence};
|
||
|
|
|
||
|
|
const QUEUE_CAPACITY: usize = 1024;
|
||
|
|
const MAX_BATCH_SIZE: usize = 64;
|
||
|
|
|
||
|
|
/// Non-blocking producer used by Wasmtime invocation threads.
|
||
|
|
pub(crate) struct InvocationPersistence {
|
||
|
|
sender: mpsc::Sender<Command>,
|
||
|
|
}
|
||
|
|
|
||
|
|
enum Command {
|
||
|
|
Record(InvocationUpdate),
|
||
|
|
Flush(oneshot::Sender<()>),
|
||
|
|
}
|
||
|
|
|
||
|
|
impl InvocationPersistence {
|
||
|
|
/// Starts the single ordered writer on the current Tokio runtime.
|
||
|
|
pub(crate) fn start(persistence: Persistence, sync: Arc<Mutex<()>>) -> Self {
|
||
|
|
let (sender, mut receiver) = mpsc::channel(QUEUE_CAPACITY);
|
||
|
|
tokio::spawn(async move {
|
||
|
|
while let Some(first) = receiver.recv().await {
|
||
|
|
let mut batch = Vec::with_capacity(MAX_BATCH_SIZE);
|
||
|
|
let mut flushes = Vec::new();
|
||
|
|
match first {
|
||
|
|
Command::Record(update) => batch.push(update),
|
||
|
|
Command::Flush(flush) => flushes.push(flush),
|
||
|
|
}
|
||
|
|
while batch.len() < MAX_BATCH_SIZE {
|
||
|
|
match receiver.try_recv() {
|
||
|
|
Ok(Command::Record(update)) => batch.push(update),
|
||
|
|
Ok(Command::Flush(flush)) => flushes.push(flush),
|
||
|
|
Err(_) => break,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if !batch.is_empty() {
|
||
|
|
let _sync = sync.lock().await;
|
||
|
|
if let Err(error) = persistence.record_invocations(batch).await {
|
||
|
|
tracing::error!(%error, "failed to persist invocation telemetry");
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for flush in flushes {
|
||
|
|
let _ = flush.send(());
|
||
|
|
}
|
||
|
|
}
|
||
|
|
});
|
||
|
|
Self { sender }
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Enqueues telemetry without extending API response latency.
|
||
|
|
pub(crate) fn record(&self, update: InvocationUpdate) {
|
||
|
|
if let Err(error) = self.sender.try_send(Command::Record(update)) {
|
||
|
|
tracing::warn!(%error, "invocation telemetry queue is unavailable; update dropped");
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Waits until all telemetry queued before this call has been attempted.
|
||
|
|
pub(crate) async fn flush(&self) {
|
||
|
|
let (sender, receiver) = oneshot::channel();
|
||
|
|
if self.sender.send(Command::Flush(sender)).await.is_ok() {
|
||
|
|
let _ = receiver.await;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|