maelstrom_util/
async.rs

1use std::{future::Future, num::NonZeroU32, sync::Mutex};
2use tokio::sync::Semaphore;
3
4pub async fn await_and_every_sec<RetT>(
5    mut fut: impl Future<Output = RetT> + Unpin,
6    mut periodic: impl FnMut(),
7) -> RetT {
8    loop {
9        tokio::select! {
10            ret = &mut fut => {
11                return ret;
12            }
13            _ = tokio::time::sleep(std::time::Duration::from_secs(1)) => {
14                periodic();
15            }
16        }
17    }
18}
19
20pub struct Pool<T> {
21    mutex: Mutex<Vec<T>>,
22    semaphore: Semaphore,
23}
24
25impl<T> Pool<T> {
26    pub fn new(slots: NonZeroU32) -> Self {
27        Self {
28            mutex: Default::default(),
29            semaphore: Semaphore::new(usize::try_from(u32::from(slots)).unwrap()),
30        }
31    }
32
33    /// Run call a closure after obtaining a slot, potentially with a pre-built item.
34    ///
35    /// The calling thread will block until there is an available slot. Once there is, the provided
36    /// closure will be called, potentially with an item returned from a previously called closure.
37    /// Upon completion, if the closure returns a new item, the item will be provided to a
38    /// subsequent closure.
39    pub async fn call_with_item<U, E, F, Fut>(&self, f: F) -> Result<U, E>
40    where
41        F: FnOnce(Option<T>) -> Fut,
42        Fut: Future<Output = Result<(T, U), E>>,
43    {
44        let _permit = self.semaphore.acquire().await.unwrap();
45        let item = self.mutex.lock().unwrap().pop();
46        f(item).await.map(|(item, ret)| {
47            self.mutex.lock().unwrap().push(item);
48            ret
49        })
50    }
51}
maelstrom_util/async.rs

maelstrom_util/
async.rs
