use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::{Duration, Instant, SystemTime};
use tokio::sync::{mpsc, oneshot};
use crate::pipeline::ProcessErrorKind;
use crate::primitives::sync::Arc;
pub fn channel() -> (DumpTrigger, DumpRx) {
let (tx, rx) = mpsc::unbounded_channel();
(DumpTrigger { tx, debounce: None }, DumpRx { rx })
}
#[derive(Debug, Default, Clone)]
pub struct DumpTriggerConfig {
debounce: Option<Duration>,
}
impl DumpTriggerConfig {
pub fn new() -> Self {
Self::default()
}
pub fn debounce(&mut self, window: Duration) {
self.debounce = Some(window);
}
pub fn debounce_window(&self) -> Option<Duration> {
self.debounce
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct DumpId(ulid::Ulid);
impl DumpId {
pub(crate) fn new() -> Self {
Self(ulid::Ulid::new())
}
pub fn timestamp(&self) -> SystemTime {
SystemTime::UNIX_EPOCH + Duration::from_millis(self.0.timestamp_ms())
}
}
impl std::fmt::Display for DumpId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
self.0.fmt(f)
}
}
impl std::str::FromStr for DumpId {
type Err = ulid::DecodeError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
Ok(Self(s.parse()?))
}
}
#[derive(Debug, Clone, Copy)]
pub(crate) enum Lookback {
Unbounded,
Window(Duration),
}
#[derive(Debug)]
pub(crate) struct DumpRequest {
pub(crate) id: DumpId,
pub(crate) triggered_at: SystemTime,
pub(crate) lookback: Lookback,
pub(crate) lookforward: Duration,
pub(crate) metadata: Vec<(String, String)>,
pub(crate) receipt_tx: oneshot::Sender<Result<DumpReceipt, DumpError>>,
}
impl DumpRequest {
pub(crate) fn elapsed_since_trigger(&self) -> Duration {
crate::primitives::time::elapsed_since(self.triggered_at)
}
}
#[derive(Debug)]
struct Debounce {
window: Duration,
last: crate::primitives::sync::Mutex<Option<(Instant, DumpId)>>,
}
#[derive(Debug, Clone)]
pub struct DumpTrigger {
tx: mpsc::UnboundedSender<DumpRequest>,
debounce: Option<Arc<Debounce>>,
}
impl DumpTrigger {
pub fn with_debounce(mut self, window: Duration) -> Self {
self.debounce = Some(Arc::new(Debounce {
window,
last: crate::primitives::sync::Mutex::new(None),
}));
self
}
pub fn dump_current_data(&self) -> DumpRun<'_> {
self.request(Lookback::Unbounded, Duration::ZERO)
}
pub fn dump_time_range(&self, lookback: Duration, lookforward: Duration) -> DumpRun<'_> {
self.request(Lookback::Window(lookback), lookforward)
}
fn request(&self, lookback: Lookback, lookforward: Duration) -> DumpRun<'_> {
let id = DumpId::new();
if let Some(debounce) = &self.debounce {
let now = Instant::now();
let mut last = debounce.last.lock().expect("debounce mutex poisoned");
match *last {
Some((at, into)) if now.duration_since(at) < debounce.window => {
return DumpRun::preempted(&self.tx, DumpError::Coalesced { into });
}
_ => *last = Some((now, id)),
}
}
let (receipt_tx, receipt_rx) = oneshot::channel();
DumpRun {
request: Some(DumpRequest {
id,
triggered_at: SystemTime::now(),
lookback,
lookforward,
metadata: Vec::new(),
receipt_tx,
}),
tx: &self.tx,
receipt_rx: Some(receipt_rx),
preempt: None,
}
}
}
#[derive(Debug)]
pub struct DumpRx {
pub(crate) rx: mpsc::UnboundedReceiver<DumpRequest>,
}
#[derive(Debug)]
pub struct DumpRun<'a> {
request: Option<DumpRequest>,
tx: &'a mpsc::UnboundedSender<DumpRequest>,
receipt_rx: Option<oneshot::Receiver<Result<DumpReceipt, DumpError>>>,
preempt: Option<DumpError>,
}
impl<'a> DumpRun<'a> {
fn preempted(tx: &'a mpsc::UnboundedSender<DumpRequest>, err: DumpError) -> Self {
DumpRun {
request: None,
tx,
receipt_rx: None,
preempt: Some(err),
}
}
pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
if let Some(req) = self.request.as_mut() {
req.metadata.push((key.into(), value.into()));
}
self
}
fn dispatch(&mut self) -> bool {
match self.request.take() {
Some(req) => self.tx.send(req).is_ok(),
None => true,
}
}
}
impl Drop for DumpRun<'_> {
fn drop(&mut self) {
let _ = self.dispatch();
}
}
impl<'a> IntoFuture for DumpRun<'a> {
type Output = Result<DumpReceipt, DumpError>;
type IntoFuture = DumpFuture;
fn into_future(mut self) -> Self::IntoFuture {
if let Some(err) = self.preempt.take() {
return DumpFuture {
inner: DumpFutureInner::Preempted(err),
};
}
let sent = self.dispatch();
let inner = match (sent, self.receipt_rx.take()) {
(true, Some(rx)) => DumpFutureInner::Waiting(rx),
_ => DumpFutureInner::Stopped,
};
DumpFuture { inner }
}
}
#[derive(Debug)]
pub struct DumpFuture {
inner: DumpFutureInner,
}
#[derive(Debug)]
enum DumpFutureInner {
Waiting(oneshot::Receiver<Result<DumpReceipt, DumpError>>),
Stopped,
Preempted(DumpError),
}
impl Future for DumpFuture {
type Output = Result<DumpReceipt, DumpError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let inner = &mut self.get_mut().inner;
match inner {
DumpFutureInner::Waiting(rx) => match Pin::new(rx).poll(cx) {
Poll::Ready(Ok(result)) => Poll::Ready(result),
Poll::Ready(Err(_)) => Poll::Ready(Err(DumpError::WorkerStopped)),
Poll::Pending => Poll::Pending,
},
DumpFutureInner::Stopped => Poll::Ready(Err(DumpError::WorkerStopped)),
DumpFutureInner::Preempted(_) => {
match std::mem::replace(inner, DumpFutureInner::Stopped) {
DumpFutureInner::Preempted(err) => Poll::Ready(Err(err)),
_ => unreachable!("matched Preempted above"),
}
}
}
}
}
#[derive(Debug)]
#[non_exhaustive]
pub struct DumpCompletion {
pub dump_id: DumpId,
pub triggered_at: SystemTime,
pub time_range: (SystemTime, SystemTime),
pub segments_processed: usize,
pub metadata: Vec<(String, String)>,
pub failed: bool,
}
#[derive(Debug)]
#[non_exhaustive]
pub struct DumpReceipt {
pub dump_id: DumpId,
pub segments_processed: usize,
pub finished_at: SystemTime,
pub time_range: (SystemTime, SystemTime),
pub manifest_key: Option<String>,
}
#[derive(Debug)]
#[non_exhaustive]
pub enum DumpError {
WorkerStopped,
Pipeline(ProcessErrorKind),
Coalesced {
into: DumpId,
},
}
impl std::fmt::Display for DumpError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::WorkerStopped => write!(f, "worker is shutting down or already stopped"),
Self::Pipeline(kind) => write!(f, "pipeline stage failed: {kind}"),
Self::Coalesced { into } => write!(f, "coalesced into dump {into}"),
}
}
}
impl std::error::Error for DumpError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::WorkerStopped => None,
Self::Pipeline(ProcessErrorKind::Io(e)) => Some(e),
Self::Pipeline(ProcessErrorKind::Transfer { source, .. }) => Some(source.as_ref()),
Self::Coalesced { .. } => None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn dispatches_on_drop_without_await() {
let (trigger, mut rx) = channel();
{
let _run = trigger
.dump_time_range(Duration::from_secs(300), Duration::from_secs(60))
.with_metadata("reason", "test")
.with_metadata("incident", "i-123");
}
let req = rx.rx.try_recv().expect("request dispatched on drop");
assert!(matches!(req.lookback, Lookback::Window(d) if d == Duration::from_secs(300)));
assert_eq!(req.lookforward, Duration::from_secs(60));
assert_eq!(
req.metadata,
vec![
("reason".to_string(), "test".to_string()),
("incident".to_string(), "i-123".to_string()),
]
);
}
#[tokio::test]
async fn dump_current_data_is_unbounded_lookback() {
let (trigger, mut rx) = channel();
trigger.dump_current_data();
let req = rx.rx.try_recv().expect("dispatched");
assert!(matches!(req.lookback, Lookback::Unbounded));
assert_eq!(req.lookforward, Duration::ZERO);
}
#[tokio::test]
async fn awaiting_dispatches_exactly_once_and_resolves_receipt() {
let (trigger, mut rx) = channel();
let run = trigger.dump_current_data();
let worker = tokio::spawn(async move {
let req = rx.rx.recv().await.expect("one request");
assert!(rx.rx.try_recv().is_err(), "no second dispatch");
let receipt = DumpReceipt {
dump_id: req.id,
segments_processed: 3,
finished_at: SystemTime::now(),
time_range: (req.triggered_at, req.triggered_at),
manifest_key: None,
};
let _ = req.receipt_tx.send(Ok(receipt));
});
let receipt = run.await.expect("receipt");
assert_eq!(receipt.segments_processed, 3);
worker.await.unwrap();
}
#[tokio::test]
async fn closed_channel_resolves_worker_stopped() {
let (trigger, rx) = channel();
drop(rx);
let err = trigger.dump_current_data().await.unwrap_err();
assert!(matches!(err, DumpError::WorkerStopped));
}
#[tokio::test]
async fn dropped_receipt_sender_resolves_worker_stopped() {
use std::future::IntoFuture;
let (trigger, mut rx) = channel();
let fut = trigger.dump_current_data().into_future();
let req = rx.rx.try_recv().expect("dispatched at into_future");
drop(req);
let err = fut.await.unwrap_err();
assert!(matches!(err, DumpError::WorkerStopped));
}
#[tokio::test]
async fn debounce_coalesces_into_the_first_dump() {
let (trigger, mut rx) = channel();
let trigger = trigger.with_debounce(Duration::from_secs(60));
let _ = trigger.dump_current_data();
let first = rx.rx.try_recv().expect("first trigger dispatched");
let err = trigger.dump_current_data().await.unwrap_err();
assert!(matches!(err, DumpError::Coalesced { into } if into == first.id));
assert!(
rx.rx.try_recv().is_err(),
"coalesced trigger must not dispatch"
);
}
#[tokio::test]
async fn debounce_dispatches_again_after_window() {
let (trigger, mut rx) = channel();
let trigger = trigger.with_debounce(Duration::from_millis(30));
let _ = trigger.dump_current_data();
let first = rx.rx.try_recv().expect("first dispatched");
tokio::time::sleep(Duration::from_millis(80)).await;
let _ = trigger.dump_current_data();
let second = rx.rx.try_recv().expect("dispatched again after window");
assert_ne!(first.id, second.id, "post-window dump gets a fresh id");
}
#[tokio::test]
async fn debounce_gate_is_shared_across_clones() {
let (trigger, mut rx) = channel();
let trigger = trigger.with_debounce(Duration::from_secs(60));
let clone = trigger.clone();
let _ = trigger.dump_current_data();
let first = rx.rx.try_recv().expect("first dispatched");
let err = clone.dump_current_data().await.unwrap_err();
assert!(matches!(err, DumpError::Coalesced { into } if into == first.id));
}
#[tokio::test]
async fn without_debounce_duplicate_triggers_both_dispatch() {
let (trigger, mut rx) = channel();
let _ = trigger.dump_current_data();
let _ = trigger.dump_current_data();
assert!(rx.rx.try_recv().is_ok(), "first dispatched");
assert!(
rx.rx.try_recv().is_ok(),
"second dispatched (no coordination)"
);
}
#[test]
fn dump_id_is_time_sorted_and_timestamp_round_trips() {
let before = SystemTime::now();
let a = DumpId::new();
std::thread::sleep(Duration::from_millis(2));
let b = DumpId::new();
let after = SystemTime::now();
assert!(a < b);
assert!(a.timestamp() >= before - Duration::from_millis(1));
assert!(b.timestamp() <= after + Duration::from_millis(1));
let parsed: DumpId = a.to_string().parse().expect("round-trip");
assert_eq!(parsed, a);
}
#[cfg(shuttle)]
mod shuttle_tests {
use super::*;
const CALLERS: usize = 4;
crate::shuttle_test! {
default;
fn shuttle_debounce_gate() {
let (trigger, mut rx) = channel();
let trigger = trigger.with_debounce(Duration::from_secs(60));
let handles: Vec<_> = (0..CALLERS)
.map(|_| {
let trigger = trigger.clone();
crate::primitives::thread::spawn(move || {
drop(trigger.dump_current_data());
})
})
.collect();
for h in handles {
h.join().unwrap();
}
let mut dispatched = 0;
while rx.rx.try_recv().is_ok() {
dispatched += 1;
}
assert_eq!(
dispatched, 1,
"exactly one trigger within the debounce window must dispatch a real request"
);
}
}
}
}