rust_zero_core/
executor.rs

1use std::{
2    fmt,
3    future::Future,
4    sync::{
5        atomic::{AtomicBool, Ordering},
6        Arc, Mutex,
7    },
8    time::{Duration, Instant},
9};
10
11use tokio::{
12    sync::{mpsc, oneshot},
13    task::JoinHandle,
14    time,
15};
16
17/// Buffers work and executes it in bounded batches by size or elapsed time.
18pub struct BatchExecutor<T> {
19    sender: mpsc::Sender<Command<T>>,
20    worker: JoinHandle<()>,
21}
22
23/// Buffers weighted items and flushes when their combined size reaches a byte limit.
24pub struct ChunkExecutor<T> {
25    sender: mpsc::Sender<ChunkCommand<T>>,
26    worker: JoinHandle<()>,
27}
28
29/// Coalesces repeated triggers into at most one delayed execution.
30pub struct DelayExecutor {
31    trigger: mpsc::Sender<()>,
32    triggered: Arc<AtomicBool>,
33    shutdown: oneshot::Sender<()>,
34    worker: JoinHandle<Result<(), String>>,
35}
36
37/// Executes at most once during each threshold window.
38pub struct LessExecutor {
39    threshold: Duration,
40    last_execution: Mutex<Option<Instant>>,
41}
42
43impl LessExecutor {
44    pub fn new(threshold: Duration) -> Self {
45        assert!(!threshold.is_zero(), "threshold must be greater than zero");
46        Self {
47            threshold,
48            last_execution: Mutex::new(None),
49        }
50    }
51
52    /// Runs `execute` when the threshold has elapsed, returning whether it ran.
53    pub fn do_or_discard<F>(&self, execute: F) -> bool
54    where
55        F: FnOnce(),
56    {
57        let now = Instant::now();
58        {
59            let mut last = self
60                .last_execution
61                .lock()
62                .unwrap_or_else(|poisoned| poisoned.into_inner());
63            if last.is_some_and(|last| now.duration_since(last) <= self.threshold) {
64                return false;
65            }
66            *last = Some(now);
67        }
68        execute();
69        true
70    }
71}
72
73impl DelayExecutor {
74    pub fn new<F, Fut, E>(delay: Duration, mut execute: F) -> Self
75    where
76        F: FnMut() -> Fut + Send + 'static,
77        Fut: Future<Output = Result<(), E>> + Send + 'static,
78        E: fmt::Display,
79    {
80        assert!(!delay.is_zero(), "delay must be greater than zero");
81        let (trigger, mut triggers) = mpsc::channel(1);
82        let triggered = Arc::new(AtomicBool::new(false));
83        let worker_triggered = Arc::clone(&triggered);
84        let (shutdown, mut stopping) = oneshot::channel();
85        let worker = tokio::spawn(async move {
86            loop {
87                tokio::select! {
88                    _ = &mut stopping => return Ok(()),
89                    trigger = triggers.recv() => {
90                        if trigger.is_none() {
91                            return Ok(());
92                        }
93                        tokio::select! {
94                            _ = &mut stopping => return Ok(()),
95                            _ = time::sleep(delay) => {}
96                        }
97                        // Match go-zero: allow a trigger made by the job to schedule another run.
98                        worker_triggered.store(false, Ordering::Release);
99                        execute().await.map_err(|error| error.to_string())?;
100                    }
101                }
102            }
103        });
104        Self {
105            trigger,
106            triggered,
107            shutdown,
108            worker,
109        }
110    }
111
112    /// Schedules the job unless a delayed execution is already pending.
113    pub fn trigger(&self) -> bool {
114        if self
115            .triggered
116            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
117            .is_err()
118        {
119            return false;
120        }
121        if self.trigger.try_send(()).is_err() {
122            self.triggered.store(false, Ordering::Release);
123            return false;
124        }
125        true
126    }
127
128    pub async fn shutdown(self, timeout: Duration) -> Result<(), DelayExecutorError> {
129        let Self {
130            trigger,
131            triggered: _,
132            shutdown,
133            mut worker,
134        } = self;
135        drop(trigger);
136        let _ = shutdown.send(());
137        match time::timeout(timeout, &mut worker).await {
138            Ok(Ok(Ok(()))) => Ok(()),
139            Ok(Ok(Err(error))) => Err(DelayExecutorError::Job(error)),
140            Ok(Err(error)) => Err(DelayExecutorError::Worker(error.to_string())),
141            Err(_) => {
142                worker.abort();
143                let _ = worker.await;
144                Err(DelayExecutorError::TimedOut(timeout))
145            }
146        }
147    }
148}
149
150#[derive(Debug, Clone, PartialEq, Eq)]
151pub enum DelayExecutorError {
152    Job(String),
153    TimedOut(Duration),
154    Worker(String),
155}
156
157impl fmt::Display for DelayExecutorError {
158    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
159        match self {
160            Self::Job(error) => write!(formatter, "delayed job failed: {error}"),
161            Self::TimedOut(timeout) => {
162                write!(formatter, "delay executor did not stop within {timeout:?}")
163            }
164            Self::Worker(error) => write!(formatter, "delay executor worker failed: {error}"),
165        }
166    }
167}
168
169impl std::error::Error for DelayExecutorError {}
170
171/// Runs one asynchronous job at a fixed interval until shutdown or the first job failure.
172///
173/// Jobs never overlap: the next interval is observed only after the current invocation
174/// completes. Shutdown is bounded and aborts a job that does not finish within the caller's
175/// deadline.
176pub struct PeriodicExecutor {
177    shutdown: oneshot::Sender<()>,
178    worker: JoinHandle<Result<(), String>>,
179}
180
181impl PeriodicExecutor {
182    pub fn new<F, Fut, E>(interval: Duration, mut execute: F) -> Self
183    where
184        F: FnMut() -> Fut + Send + 'static,
185        Fut: Future<Output = Result<(), E>> + Send + 'static,
186        E: fmt::Display,
187    {
188        assert!(!interval.is_zero(), "interval must be greater than zero");
189        let (shutdown, mut stopping) = oneshot::channel();
190        let worker = tokio::spawn(async move {
191            let start = time::Instant::now() + interval;
192            let mut ticker = time::interval_at(start, interval);
193            ticker.set_missed_tick_behavior(time::MissedTickBehavior::Delay);
194
195            loop {
196                tokio::select! {
197                    _ = &mut stopping => return Ok(()),
198                    _ = ticker.tick() => {
199                        execute().await.map_err(|error| error.to_string())?;
200                    }
201                }
202            }
203        });
204        Self { shutdown, worker }
205    }
206
207    /// Requests shutdown and waits no longer than `timeout` for an active job to finish.
208    pub async fn shutdown(self, timeout: Duration) -> Result<(), PeriodicExecutorError> {
209        let Self {
210            shutdown,
211            mut worker,
212        } = self;
213        let _ = shutdown.send(());
214
215        match time::timeout(timeout, &mut worker).await {
216            Ok(Ok(Ok(()))) => Ok(()),
217            Ok(Ok(Err(error))) => Err(PeriodicExecutorError::Job(error)),
218            Ok(Err(error)) => Err(PeriodicExecutorError::Worker(error.to_string())),
219            Err(_) => {
220                worker.abort();
221                let _ = worker.await;
222                Err(PeriodicExecutorError::TimedOut(timeout))
223            }
224        }
225    }
226}
227
228#[derive(Debug, Clone, PartialEq, Eq)]
229pub enum PeriodicExecutorError {
230    Job(String),
231    TimedOut(Duration),
232    Worker(String),
233}
234
235impl fmt::Display for PeriodicExecutorError {
236    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
237        match self {
238            Self::Job(error) => write!(formatter, "periodic job failed: {error}"),
239            Self::TimedOut(timeout) => {
240                write!(
241                    formatter,
242                    "periodic executor did not stop within {timeout:?}"
243                )
244            }
245            Self::Worker(error) => write!(formatter, "periodic executor worker failed: {error}"),
246        }
247    }
248}
249
250impl std::error::Error for PeriodicExecutorError {}
251
252impl<T> BatchExecutor<T>
253where
254    T: Send + 'static,
255{
256    pub fn new<F, Fut>(max_batch_size: usize, flush_interval: Duration, execute: F) -> Self
257    where
258        F: FnMut(Vec<T>) -> Fut + Send + 'static,
259        Fut: Future<Output = ()> + Send + 'static,
260    {
261        assert!(
262            max_batch_size > 0,
263            "maximum batch size must be greater than zero"
264        );
265        assert!(
266            !flush_interval.is_zero(),
267            "flush interval must be greater than zero"
268        );
269
270        let (sender, receiver) = mpsc::channel(max_batch_size.saturating_mul(2).max(1));
271        let worker = tokio::spawn(run_worker(
272            receiver,
273            max_batch_size,
274            flush_interval,
275            execute,
276        ));
277        Self { sender, worker }
278    }
279
280    /// Queues one item, applying backpressure when the input buffer is full.
281    pub async fn push(&self, value: T) -> Result<(), BatchExecutorError> {
282        self.sender
283            .send(Command::Item(value))
284            .await
285            .map_err(|_| BatchExecutorError::Closed)
286    }
287
288    /// Executes all items accepted before this call and waits for completion.
289    pub async fn flush(&self) -> Result<(), BatchExecutorError> {
290        let (sender, receiver) = oneshot::channel();
291        self.sender
292            .send(Command::Flush(sender))
293            .await
294            .map_err(|_| BatchExecutorError::Closed)?;
295        receiver.await.map_err(|_| BatchExecutorError::Closed)
296    }
297
298    /// Flushes pending work and stops the worker.
299    pub async fn shutdown(self) -> Result<(), BatchExecutorError> {
300        let (sender, receiver) = oneshot::channel();
301        self.sender
302            .send(Command::Shutdown(sender))
303            .await
304            .map_err(|_| BatchExecutorError::Closed)?;
305        receiver.await.map_err(|_| BatchExecutorError::Closed)?;
306        self.worker
307            .await
308            .map_err(|error| BatchExecutorError::Worker(error.to_string()))
309    }
310}
311
312impl<T> ChunkExecutor<T>
313where
314    T: Send + 'static,
315{
316    pub fn new<F, Fut>(max_chunk_bytes: usize, flush_interval: Duration, execute: F) -> Self
317    where
318        F: FnMut(Vec<T>) -> Fut + Send + 'static,
319        Fut: Future<Output = ()> + Send + 'static,
320    {
321        assert!(max_chunk_bytes > 0, "chunk size must be greater than zero");
322        assert!(
323            !flush_interval.is_zero(),
324            "flush interval must be greater than zero"
325        );
326
327        let (sender, receiver) = mpsc::channel(128);
328        let worker = tokio::spawn(run_chunk_worker(
329            receiver,
330            max_chunk_bytes,
331            flush_interval,
332            execute,
333        ));
334        Self { sender, worker }
335    }
336
337    /// Queues an item with its accounting size, applying backpressure when full.
338    pub async fn push(&self, value: T, size: usize) -> Result<(), BatchExecutorError> {
339        self.sender
340            .send(ChunkCommand::Item { value, size })
341            .await
342            .map_err(|_| BatchExecutorError::Closed)
343    }
344
345    /// Executes all items accepted before this call and waits for completion.
346    pub async fn flush(&self) -> Result<(), BatchExecutorError> {
347        let (sender, receiver) = oneshot::channel();
348        self.sender
349            .send(ChunkCommand::Flush(sender))
350            .await
351            .map_err(|_| BatchExecutorError::Closed)?;
352        receiver.await.map_err(|_| BatchExecutorError::Closed)
353    }
354
355    /// Flushes pending work and stops the worker.
356    pub async fn shutdown(self) -> Result<(), BatchExecutorError> {
357        let (sender, receiver) = oneshot::channel();
358        self.sender
359            .send(ChunkCommand::Shutdown(sender))
360            .await
361            .map_err(|_| BatchExecutorError::Closed)?;
362        receiver.await.map_err(|_| BatchExecutorError::Closed)?;
363        self.worker
364            .await
365            .map_err(|error| BatchExecutorError::Worker(error.to_string()))
366    }
367}
368
369enum Command<T> {
370    Item(T),
371    Flush(oneshot::Sender<()>),
372    Shutdown(oneshot::Sender<()>),
373}
374
375enum ChunkCommand<T> {
376    Item { value: T, size: usize },
377    Flush(oneshot::Sender<()>),
378    Shutdown(oneshot::Sender<()>),
379}
380
381async fn run_worker<T, F, Fut>(
382    mut receiver: mpsc::Receiver<Command<T>>,
383    max_batch_size: usize,
384    flush_interval: Duration,
385    mut execute: F,
386) where
387    T: Send + 'static,
388    F: FnMut(Vec<T>) -> Fut + Send + 'static,
389    Fut: Future<Output = ()> + Send + 'static,
390{
391    let mut batch = Vec::with_capacity(max_batch_size);
392    let mut ticker = time::interval(flush_interval);
393    ticker.set_missed_tick_behavior(time::MissedTickBehavior::Delay);
394    ticker.tick().await;
395
396    loop {
397        tokio::select! {
398            command = receiver.recv() => {
399                match command {
400                    Some(Command::Item(value)) => {
401                        batch.push(value);
402                        if batch.len() >= max_batch_size {
403                            execute(std::mem::take(&mut batch)).await;
404                            batch.reserve(max_batch_size);
405                        }
406                    }
407                    Some(Command::Flush(done)) => {
408                        flush_batch(&mut batch, &mut execute).await;
409                        let _ = done.send(());
410                    }
411                    Some(Command::Shutdown(done)) => {
412                        flush_batch(&mut batch, &mut execute).await;
413                        let _ = done.send(());
414                        return;
415                    }
416                    None => {
417                        flush_batch(&mut batch, &mut execute).await;
418                        return;
419                    }
420                }
421            }
422            _ = ticker.tick() => {
423                flush_batch(&mut batch, &mut execute).await;
424            }
425        }
426    }
427}
428
429async fn run_chunk_worker<T, F, Fut>(
430    mut receiver: mpsc::Receiver<ChunkCommand<T>>,
431    max_chunk_bytes: usize,
432    flush_interval: Duration,
433    mut execute: F,
434) where
435    T: Send + 'static,
436    F: FnMut(Vec<T>) -> Fut + Send + 'static,
437    Fut: Future<Output = ()> + Send + 'static,
438{
439    let mut chunk = Vec::new();
440    let mut chunk_bytes = 0usize;
441    let mut ticker = time::interval(flush_interval);
442    ticker.set_missed_tick_behavior(time::MissedTickBehavior::Delay);
443    ticker.tick().await;
444
445    loop {
446        tokio::select! {
447            command = receiver.recv() => {
448                match command {
449                    Some(ChunkCommand::Item { value, size }) => {
450                        chunk.push(value);
451                        chunk_bytes = chunk_bytes.saturating_add(size);
452                        if chunk_bytes >= max_chunk_bytes {
453                            execute(std::mem::take(&mut chunk)).await;
454                            chunk_bytes = 0;
455                        }
456                    }
457                    Some(ChunkCommand::Flush(done)) => {
458                        flush_batch(&mut chunk, &mut execute).await;
459                        chunk_bytes = 0;
460                        let _ = done.send(());
461                    }
462                    Some(ChunkCommand::Shutdown(done)) => {
463                        flush_batch(&mut chunk, &mut execute).await;
464                        let _ = done.send(());
465                        return;
466                    }
467                    None => {
468                        flush_batch(&mut chunk, &mut execute).await;
469                        return;
470                    }
471                }
472            }
473            _ = ticker.tick() => {
474                flush_batch(&mut chunk, &mut execute).await;
475                chunk_bytes = 0;
476            }
477        }
478    }
479}
480
481async fn flush_batch<T, F, Fut>(batch: &mut Vec<T>, execute: &mut F)
482where
483    F: FnMut(Vec<T>) -> Fut,
484    Fut: Future<Output = ()>,
485{
486    if !batch.is_empty() {
487        execute(std::mem::take(batch)).await;
488    }
489}
490
491#[derive(Debug, Clone, PartialEq, Eq)]
492pub enum BatchExecutorError {
493    Closed,
494    Worker(String),
495}
496
497impl fmt::Display for BatchExecutorError {
498    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
499        match self {
500            Self::Closed => formatter.write_str("batch executor is closed"),
501            Self::Worker(error) => write!(formatter, "batch executor worker failed: {error}"),
502        }
503    }
504}
505
506impl std::error::Error for BatchExecutorError {}
507
508#[cfg(test)]
509mod tests {
510    use super::*;
511    use tokio::sync::mpsc;
512
513    #[tokio::test]
514    async fn flushes_when_the_batch_reaches_its_limit() {
515        let (batches, mut received) = mpsc::unbounded_channel();
516        let executor = BatchExecutor::new(3, Duration::from_secs(60), move |batch| {
517            let batches = batches.clone();
518            async move {
519                batches.send(batch).unwrap();
520            }
521        });
522
523        executor.push(1).await.unwrap();
524        executor.push(2).await.unwrap();
525        executor.push(3).await.unwrap();
526
527        assert_eq!(received.recv().await.unwrap(), vec![1, 2, 3]);
528        executor.shutdown().await.unwrap();
529    }
530
531    #[tokio::test]
532    async fn explicit_flush_and_shutdown_do_not_lose_work() {
533        let (batches, mut received) = mpsc::unbounded_channel();
534        let executor = BatchExecutor::new(10, Duration::from_secs(60), move |batch| {
535            let batches = batches.clone();
536            async move {
537                batches.send(batch).unwrap();
538            }
539        });
540
541        executor.push("first").await.unwrap();
542        executor.flush().await.unwrap();
543        assert_eq!(received.recv().await.unwrap(), vec!["first"]);
544
545        executor.push("second").await.unwrap();
546        executor.shutdown().await.unwrap();
547        assert_eq!(received.recv().await.unwrap(), vec!["second"]);
548    }
549
550    #[tokio::test]
551    async fn chunk_executor_flushes_on_combined_size() {
552        let (chunks, mut received) = mpsc::unbounded_channel();
553        let executor = ChunkExecutor::new(10, Duration::from_secs(60), move |chunk| {
554            let chunks = chunks.clone();
555            async move {
556                chunks.send(chunk).unwrap();
557            }
558        });
559
560        executor.push("small", 4).await.unwrap();
561        executor.push("large", 6).await.unwrap();
562        assert_eq!(received.recv().await.unwrap(), vec!["small", "large"]);
563        executor.shutdown().await.unwrap();
564    }
565
566    #[tokio::test]
567    async fn chunk_executor_flushes_on_interval_and_shutdown() {
568        let (chunks, mut received) = mpsc::unbounded_channel();
569        let executor = ChunkExecutor::new(100, Duration::from_millis(5), move |chunk| {
570            let chunks = chunks.clone();
571            async move {
572                chunks.send(chunk).unwrap();
573            }
574        });
575
576        executor.push(1, 1).await.unwrap();
577        assert_eq!(received.recv().await.unwrap(), vec![1]);
578        executor.push(2, 1).await.unwrap();
579        executor.shutdown().await.unwrap();
580        assert_eq!(received.recv().await.unwrap(), vec![2]);
581    }
582
583    #[tokio::test]
584    async fn delay_executor_coalesces_pending_triggers_and_can_run_again() {
585        let (runs, mut received) = mpsc::unbounded_channel();
586        let executor = DelayExecutor::new(Duration::from_millis(5), move || {
587            let runs = runs.clone();
588            async move {
589                runs.send(()).unwrap();
590                Ok::<_, &'static str>(())
591            }
592        });
593
594        assert!(executor.trigger());
595        assert!(!executor.trigger());
596        time::timeout(Duration::from_secs(1), received.recv())
597            .await
598            .unwrap()
599            .unwrap();
600        assert!(executor.trigger());
601        time::timeout(Duration::from_secs(1), received.recv())
602            .await
603            .unwrap()
604            .unwrap();
605        executor.shutdown(Duration::from_secs(1)).await.unwrap();
606    }
607
608    #[tokio::test]
609    async fn delay_executor_reports_job_failure() {
610        let executor = DelayExecutor::new(Duration::from_millis(1), || async {
611            Err::<(), _>("write failed")
612        });
613        assert!(executor.trigger());
614        time::sleep(Duration::from_millis(10)).await;
615        assert_eq!(
616            executor.shutdown(Duration::from_secs(1)).await,
617            Err(DelayExecutorError::Job("write failed".to_owned()))
618        );
619    }
620
621    #[test]
622    fn less_executor_discards_calls_inside_the_threshold() {
623        let executor = LessExecutor::new(Duration::from_millis(10));
624        let mut runs = 0;
625        assert!(executor.do_or_discard(|| runs += 1));
626        assert!(!executor.do_or_discard(|| runs += 1));
627        std::thread::sleep(Duration::from_millis(15));
628        assert!(executor.do_or_discard(|| runs += 1));
629        assert_eq!(runs, 2);
630    }
631
632    #[tokio::test]
633    async fn periodic_executor_runs_repeatedly_and_stops() {
634        let (runs, mut received) = mpsc::unbounded_channel();
635        let executor = PeriodicExecutor::new(Duration::from_millis(5), move || {
636            let runs = runs.clone();
637            async move {
638                runs.send(()).unwrap();
639                Ok::<_, &'static str>(())
640            }
641        });
642
643        time::timeout(Duration::from_secs(1), received.recv())
644            .await
645            .unwrap()
646            .unwrap();
647        time::timeout(Duration::from_secs(1), received.recv())
648            .await
649            .unwrap()
650            .unwrap();
651        executor.shutdown(Duration::from_secs(1)).await.unwrap();
652    }
653
654    #[tokio::test]
655    async fn periodic_executor_reports_job_failure() {
656        let executor = PeriodicExecutor::new(Duration::from_millis(1), || async {
657            Err::<(), _>("backend unavailable")
658        });
659        time::sleep(Duration::from_millis(10)).await;
660
661        assert_eq!(
662            executor.shutdown(Duration::from_secs(1)).await,
663            Err(PeriodicExecutorError::Job("backend unavailable".to_owned()))
664        );
665    }
666
667    #[tokio::test]
668    async fn periodic_executor_bounds_slow_shutdown() {
669        let (started, start_received) = oneshot::channel();
670        let mut started = Some(started);
671        let executor = PeriodicExecutor::new(Duration::from_millis(1), move || {
672            if let Some(started) = started.take() {
673                let _ = started.send(());
674            }
675            async {
676                std::future::pending::<()>().await;
677                Ok::<_, &'static str>(())
678            }
679        });
680        time::timeout(Duration::from_secs(1), start_received)
681            .await
682            .unwrap()
683            .unwrap();
684
685        assert_eq!(
686            executor.shutdown(Duration::from_millis(5)).await,
687            Err(PeriodicExecutorError::TimedOut(Duration::from_millis(5)))
688        );
689    }
690}
rust_zero_core/executor.rs

rust_zero_core/
executor.rs
