mocra_core/queue/
batcher.rs

1use std::sync::Arc;
2use std::time::Duration;
3use tokio::sync::Semaphore;
4use tokio::sync::mpsc::Receiver;
5
6pub struct Batcher;
7
8impl Batcher {
9    pub async fn run<T, F, Fut>(
10        rx: &mut Receiver<T>,
11        batch_size: usize,
12        interval_ms: u64,
13        semaphore: Arc<Semaphore>,
14        processor: F,
15    ) where
16        T: Send + 'static,
17        F: Fn(Vec<T>) -> Fut + Send + Sync + 'static + Clone,
18        Fut: std::future::Future<Output = ()> + Send + 'static,
19    {
20        let mut batch = Vec::with_capacity(batch_size);
21        let mut interval = tokio::time::interval(Duration::from_millis(interval_ms));
22        loop {
23            tokio::select! {
24                res = rx.recv() => {
25                    match res {
26                        Some(item) => {
27                            // println!("DEBUG: Batcher received item");
28                            batch.push(item);
29                            while batch.len() < batch_size {
30                                match rx.try_recv() {
31                                    Ok(next) => batch.push(next),
32                                    Err(_) => break,
33                                }
34                            }
35
36                            if batch.len() >= batch_size {
37                                let items = std::mem::replace(&mut batch, Vec::with_capacity(batch_size));
38                                let processor = processor.clone();
39                                let semaphore = semaphore.clone();
40
41                                tokio::spawn(async move {
42                                    if let Ok(_permit) = semaphore.acquire_owned().await {
43                                        processor(items).await;
44                                    }
45                                });
46                            }
47                        }
48                        None => {
49                            if !batch.is_empty() {
50                                let items = std::mem::take(&mut batch);
51                                let processor = processor.clone();
52                                let semaphore = semaphore.clone();
53
54                                tokio::spawn(async move {
55                                    if let Ok(_permit) = semaphore.acquire_owned().await {
56                                        processor(items).await;
57                                    }
58                                });
59                            }
60                            break;
61                        }
62                    }
63                }
64                _ = interval.tick() => {
65                    if !batch.is_empty() {
66                         let items = std::mem::replace(&mut batch, Vec::with_capacity(batch_size));
67                         let processor = processor.clone();
68                         let semaphore = semaphore.clone();
69
70                         tokio::spawn(async move {
71                             if let Ok(_permit) = semaphore.acquire_owned().await {
72                                 processor(items).await;
73                             }
74                         });
75                    }
76                }
77            }
78        }
79    }
80}
mocra_core/queue/batcher.rs

mocra_core/queue/
batcher.rs
