async_repeater/
repeater.rs

1use crate::{
2    Delay, RepeaterEntry,
3    handle::{Message, RepeaterHandle},
4};
5use futures_core::ready;
6use std::{
7    collections::HashMap,
8    future::Future,
9    pin::Pin,
10    task::{Context, Poll},
11    time::{Duration, SystemTime},
12};
13use tokio::sync::mpsc::{Receiver, channel};
14use tokio_stream::{self, Stream, StreamExt};
15use tokio_util::time::{DelayQueue, delay_queue};
16
17/// Repeats structs after user defined Duration.
18pub struct Repeater<E>
19where
20    E: RepeaterEntry + Clone + Unpin + Send,
21{
22    entries: HashMap<E::Key, (E, delay_queue::Key)>,
23    queue: DelayQueue<E::Key>,
24}
25
26impl<E> Repeater<E>
27where
28    E: RepeaterEntry + Clone + Unpin + Send + Sync,
29{
30    /// Create a new repeater with given capacity
31    pub fn with_capacity(capacity: usize) -> Self {
32        Self {
33            entries: HashMap::with_capacity(capacity),
34            queue: DelayQueue::with_capacity(capacity),
35        }
36    }
37
38    /// Insert entry into repeater
39    ///
40    /// If entry with same key already exists, it will be replaced.
41    /// Replacement will also reset the repetition.
42    ///
43    /// Insertion respects the delay() call
44    pub fn insert(&mut self, e: E) {
45        // if there is a delay before the first run
46        let interval = match e.delay() {
47            Delay::Relative(dur) => dur,
48            Delay::Absolute(inst) => inst.duration_since(SystemTime::now()).unwrap_or_else(|_| e.interval()),
49            Delay::None => Duration::default(),
50        };
51
52        match self.entries.get_mut(&e.key()) {
53            Some((current_item, queue_key)) => {
54                // if an item with the same key already exists,
55                // then replace that item with the new one and set the timer to the new value as well
56                self.queue.reset(queue_key, interval);
57                *current_item = e;
58            }
59            _ => {
60                // if the item is new, insert it and the time at which to start
61                let queue_key = self.queue.insert(e.key(), interval);
62                self.entries.insert(e.key(), (e.clone(), queue_key));
63            }
64        }
65    }
66
67    /// Remove entry with key from repeater
68    pub fn remove(&mut self, key: &E::Key) {
69        let Some((_, key)) = self.entries.remove(key) else {
70            return;
71        };
72        self.queue.remove(&key);
73    }
74
75    /// Clear the repeater. This removes all entries.
76    pub fn clear(&mut self) {
77        self.queue.clear();
78        self.entries.clear();
79    }
80
81    /// Number of entries in the repeater
82    pub fn len(&self) -> usize {
83        self.queue.len()
84    }
85
86    /// Return true if repeater is empty
87    pub fn is_empty(&self) -> bool {
88        self.queue.is_empty()
89    }
90
91    /// Starts the repeater in a background task
92    ///
93    /// This consumes the repeater and returns a [`RepeaterHandle`] which allows communication with
94    /// the background task.
95    pub fn run_with_async_callback<F, Fut>(self, callback: F) -> RepeaterHandle<E>
96    where
97        F: FnOnce(E, RepeaterHandle<E>) -> Fut + Send + 'static + Clone,
98        Fut: Future<Output = ()> + Send + 'static,
99    {
100        let (tx, rx) = channel(self.queue.capacity());
101        let handle = RepeaterHandle::new(tx);
102
103        tokio::spawn(self.handle_messages(rx, handle.clone(), callback));
104
105        handle
106    }
107
108    /// Handle communication from the [`RepeaterHandle`].
109    ///
110    /// It's mainly a separate function to decrease indention inside [`Self::run_with_async_callback`].
111    /// This is run in a separate task
112    async fn handle_messages<F, Fut>(mut self, mut rx: Receiver<Message<E>>, handle: RepeaterHandle<E>, callback: F)
113    where
114        F: FnOnce(E, RepeaterHandle<E>) -> Fut + Send + 'static + Clone,
115        Fut: Future<Output = ()> + Send + 'static,
116    {
117        loop {
118            let cb = callback.clone();
119            let handle = handle.clone();
120
121            tokio::select! {
122                Some(message) = rx.recv() => {
123                    match message {
124                        Message::Insert(entry) => self.insert(entry.clone()),
125                        Message::Remove(key) => self.remove(&key),
126                        Message::Clear => { self.clear()},
127                        Message::Stop => { break },
128                        Message::Len(reply_sender) => { let _ = reply_sender.send(self.len()); }
129                        Message::IsEmtpy(reply_sender) => { let _ = reply_sender.send(self.is_empty()); }
130                    }
131                }
132                Some(entry) = self.next() => {
133                    tokio::spawn((cb)(entry, handle));
134                }
135            }
136        }
137    }
138
139    /// Poll for a new item.
140    ///
141    /// If an item is available it will also be re-inserted.
142    /// Re-insertions ignore the delay and will only respect the interval.
143    fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<E>> {
144        let entry_id: Option<E::Key> = ready!(self.queue.poll_expired(cx)).map(delay_queue::Expired::into_inner);
145        if let Some(entry_id) = entry_id {
146            let (entry, queue_key) = self.entries.get_mut(&entry_id).unwrap();
147
148            let new_queue_key = self.queue.insert(entry.key(), entry.interval());
149            *queue_key = new_queue_key;
150
151            return Poll::Ready(Some(entry.clone()));
152        }
153
154        Poll::Pending
155    }
156}
157
158impl<E> Stream for Repeater<E>
159where
160    E: RepeaterEntry + Clone + Unpin + Send + Sync,
161{
162    type Item = E;
163
164    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<<Self as Stream>::Item>> {
165        Repeater::poll_next(self.get_mut(), cx)
166    }
167}
async_repeater/repeater.rs

async_repeater/
repeater.rs
