tenant.openagents/omega

1use std::{
2    collections::VecDeque,
3    fmt,
4    iter::FusedIterator,
5    sync::{Arc, atomic::AtomicUsize},
6};
7
8use rand::{Rng, SeedableRng, rngs::SmallRng};
9
10use crate::Priority;
11
12struct PriorityQueues<T> {
13    high_priority: VecDeque<T>,
14    medium_priority: VecDeque<T>,
15    low_priority: VecDeque<T>,
16}
17
18impl<T> PriorityQueues<T> {
19    fn is_empty(&self) -> bool {
20        self.high_priority.is_empty()
21            && self.medium_priority.is_empty()
22            && self.low_priority.is_empty()
23    }
24}
25
26struct PriorityQueueState<T> {
27    queues: parking_lot::Mutex<PriorityQueues<T>>,
28    condvar: parking_lot::Condvar,
29    receiver_count: AtomicUsize,
30    sender_count: AtomicUsize,
31}
32
33impl<T> PriorityQueueState<T> {
34    fn send(&self, priority: Priority, item: T) -> Result<(), SendError<T>> {
35        if self
36            .receiver_count
37            .load(std::sync::atomic::Ordering::Relaxed)
38            == 0
39        {
40            return Err(SendError(item));
41        }
42
43        let mut queues = self.queues.lock();
44        Self::push(&mut queues, priority, item);
45        self.condvar.notify_one();
46        Ok(())
47    }
48
49    fn spin_send(&self, priority: Priority, item: T) -> Result<(), SendError<T>> {
50        if self
51            .receiver_count
52            .load(std::sync::atomic::Ordering::Relaxed)
53            == 0
54        {
55            return Err(SendError(item));
56        }
57
58        let mut queues = loop {
59            if let Some(guard) = self.queues.try_lock() {
60                break guard;
61            }
62            std::hint::spin_loop();
63        };
64        Self::push(&mut queues, priority, item);
65        self.condvar.notify_one();
66        Ok(())
67    }
68
69    fn push(queues: &mut PriorityQueues<T>, priority: Priority, item: T) {
70        match priority {
71            Priority::RealtimeAudio => unreachable!(
72                "Realtime audio priority runs on a dedicated thread and is never queued"
73            ),
74            Priority::High => queues.high_priority.push_back(item),
75            Priority::Medium => queues.medium_priority.push_back(item),
76            Priority::Low => queues.low_priority.push_back(item),
77        };
78    }
79
80    fn recv<'a>(&'a self) -> Result<parking_lot::MutexGuard<'a, PriorityQueues<T>>, RecvError> {
81        let mut queues = self.queues.lock();
82
83        let sender_count = self.sender_count.load(std::sync::atomic::Ordering::Relaxed);
84        if queues.is_empty() && sender_count == 0 {
85            return Err(crate::queue::RecvError);
86        }
87
88        while queues.is_empty() {
89            self.condvar.wait(&mut queues);
90        }
91
92        Ok(queues)
93    }
94
95    fn try_recv<'a>(
96        &'a self,
97    ) -> Result<Option<parking_lot::MutexGuard<'a, PriorityQueues<T>>>, RecvError> {
98        let mut queues = self.queues.lock();
99
100        let sender_count = self.sender_count.load(std::sync::atomic::Ordering::Relaxed);
101        if queues.is_empty() && sender_count == 0 {
102            return Err(crate::queue::RecvError);
103        }
104
105        if queues.is_empty() {
106            Ok(None)
107        } else {
108            Ok(Some(queues))
109        }
110    }
111
112    fn spin_try_recv<'a>(
113        &'a self,
114    ) -> Result<Option<parking_lot::MutexGuard<'a, PriorityQueues<T>>>, RecvError> {
115        let queues = loop {
116            if let Some(guard) = self.queues.try_lock() {
117                break guard;
118            }
119            std::hint::spin_loop();
120        };
121
122        let sender_count = self.sender_count.load(std::sync::atomic::Ordering::Relaxed);
123        if queues.is_empty() && sender_count == 0 {
124            return Err(crate::queue::RecvError);
125        }
126
127        if queues.is_empty() {
128            Ok(None)
129        } else {
130            Ok(Some(queues))
131        }
132    }
133}
134
135#[doc(hidden)]
136pub struct PriorityQueueSender<T> {
137    state: Arc<PriorityQueueState<T>>,
138}
139
140impl<T> PriorityQueueSender<T> {
141    fn new(state: Arc<PriorityQueueState<T>>) -> Self {
142        Self { state }
143    }
144
145    pub fn send(&self, priority: Priority, item: T) -> Result<(), SendError<T>> {
146        self.state.send(priority, item)?;
147        Ok(())
148    }
149
150    pub fn spin_send(&self, priority: Priority, item: T) -> Result<(), SendError<T>> {
151        self.state.spin_send(priority, item)?;
152        Ok(())
153    }
154}
155
156impl<T> Drop for PriorityQueueSender<T> {
157    fn drop(&mut self) {
158        self.state
159            .sender_count
160            .fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
161    }
162}
163
164#[doc(hidden)]
165pub struct PriorityQueueReceiver<T> {
166    state: Arc<PriorityQueueState<T>>,
167    rand: SmallRng,
168    disconnected: bool,
169}
170
171impl<T> Clone for PriorityQueueReceiver<T> {
172    fn clone(&self) -> Self {
173        self.state
174            .receiver_count
175            .fetch_add(1, std::sync::atomic::Ordering::AcqRel);
176        Self {
177            state: Arc::clone(&self.state),
178            rand: SmallRng::seed_from_u64(0),
179            disconnected: self.disconnected,
180        }
181    }
182}
183
184#[doc(hidden)]
185pub struct SendError<T>(pub T);
186
187impl<T: fmt::Debug> fmt::Debug for SendError<T> {
188    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
189        f.debug_tuple("SendError").field(&self.0).finish()
190    }
191}
192
193#[derive(Debug)]
194#[doc(hidden)]
195pub struct RecvError;
196
197#[allow(dead_code)]
198impl<T> PriorityQueueReceiver<T> {
199    pub fn new() -> (PriorityQueueSender<T>, Self) {
200        let state = PriorityQueueState {
201            queues: parking_lot::Mutex::new(PriorityQueues {
202                high_priority: VecDeque::new(),
203                medium_priority: VecDeque::new(),
204                low_priority: VecDeque::new(),
205            }),
206            condvar: parking_lot::Condvar::new(),
207            receiver_count: AtomicUsize::new(1),
208            sender_count: AtomicUsize::new(1),
209        };
210        let state = Arc::new(state);
211
212        let sender = PriorityQueueSender::new(Arc::clone(&state));
213
214        let receiver = PriorityQueueReceiver {
215            state,
216            rand: SmallRng::seed_from_u64(0),
217            disconnected: false,
218        };
219
220        (sender, receiver)
221    }
222
223    /// Returns whether the queue currently contains no elements.
224    pub fn is_empty(&self) -> bool {
225        self.state.queues.lock().is_empty()
226    }
227
228    /// Tries to pop one element from the priority queue without blocking.
229    ///
230    /// This will early return if there are no elements in the queue.
231    ///
232    /// This method is best suited if you only intend to pop one element, for better performance
233    /// on large queues see [`Self::try_iter`]
234    ///
235    /// # Errors
236    ///
237    /// If the sender was dropped
238    pub fn try_pop(&mut self) -> Result<Option<T>, RecvError> {
239        self.pop_inner(false)
240    }
241
242    pub fn spin_try_pop(&mut self) -> Result<Option<T>, RecvError> {
243        use Priority as P;
244
245        let Some(mut queues) = self.state.spin_try_recv()? else {
246            return Ok(None);
247        };
248
249        let high = P::High.weight() * !queues.high_priority.is_empty() as u32;
250        let medium = P::Medium.weight() * !queues.medium_priority.is_empty() as u32;
251        let low = P::Low.weight() * !queues.low_priority.is_empty() as u32;
252        let mut mass = high + medium + low;
253
254        if !queues.high_priority.is_empty() {
255            let flip = self.rand.random_ratio(P::High.weight(), mass);
256            if flip {
257                return Ok(queues.high_priority.pop_front());
258            }
259            mass -= P::High.weight();
260        }
261
262        if !queues.medium_priority.is_empty() {
263            let flip = self.rand.random_ratio(P::Medium.weight(), mass);
264            if flip {
265                return Ok(queues.medium_priority.pop_front());
266            }
267            mass -= P::Medium.weight();
268        }
269
270        if !queues.low_priority.is_empty() {
271            let flip = self.rand.random_ratio(P::Low.weight(), mass);
272            if flip {
273                return Ok(queues.low_priority.pop_front());
274            }
275        }
276
277        Ok(None)
278    }
279
280    /// Pops an element from the priority queue blocking if necessary.
281    ///
282    /// This method is best suited if you only intend to pop one element, for better performance
283    /// on large queues see [`Self::iter``]
284    ///
285    /// # Errors
286    ///
287    /// If the sender was dropped
288    pub fn pop(&mut self) -> Result<T, RecvError> {
289        self.pop_inner(true).map(|e| e.unwrap())
290    }
291
292    /// Returns an iterator over the elements of the queue
293    /// this iterator will end when all elements have been consumed and will not wait for new ones.
294    pub fn try_iter(self) -> TryIter<T> {
295        TryIter {
296            receiver: self,
297            ended: false,
298        }
299    }
300
301    /// Returns an iterator over the elements of the queue
302    /// this iterator will wait for new elements if the queue is empty.
303    pub fn iter(self) -> Iter<T> {
304        Iter(self)
305    }
306
307    #[inline(always)]
308    // algorithm is the loaded die from biased coin from
309    // https://www.keithschwarz.com/darts-dice-coins/
310    fn pop_inner(&mut self, block: bool) -> Result<Option<T>, RecvError> {
311        use Priority as P;
312
313        let mut queues = if !block {
314            let Some(queues) = self.state.try_recv()? else {
315                return Ok(None);
316            };
317            queues
318        } else {
319            self.state.recv()?
320        };
321
322        let high = P::High.weight() * !queues.high_priority.is_empty() as u32;
323        let medium = P::Medium.weight() * !queues.medium_priority.is_empty() as u32;
324        let low = P::Low.weight() * !queues.low_priority.is_empty() as u32;
325        let mut mass = high + medium + low; //%
326
327        if !queues.high_priority.is_empty() {
328            let flip = self.rand.random_ratio(P::High.weight(), mass);
329            if flip {
330                return Ok(queues.high_priority.pop_front());
331            }
332            mass -= P::High.weight();
333        }
334
335        if !queues.medium_priority.is_empty() {
336            let flip = self.rand.random_ratio(P::Medium.weight(), mass);
337            if flip {
338                return Ok(queues.medium_priority.pop_front());
339            }
340            mass -= P::Medium.weight();
341        }
342
343        if !queues.low_priority.is_empty() {
344            let flip = self.rand.random_ratio(P::Low.weight(), mass);
345            if flip {
346                return Ok(queues.low_priority.pop_front());
347            }
348        }
349
350        Ok(None)
351    }
352}
353
354impl<T> Drop for PriorityQueueReceiver<T> {
355    fn drop(&mut self) {
356        self.state
357            .receiver_count
358            .fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
359    }
360}
361
362#[doc(hidden)]
363pub struct Iter<T>(PriorityQueueReceiver<T>);
364impl<T> Iterator for Iter<T> {
365    type Item = T;
366
367    fn next(&mut self) -> Option<Self::Item> {
368        self.0.pop().ok()
369    }
370}
371impl<T> FusedIterator for Iter<T> {}
372
373#[doc(hidden)]
374pub struct TryIter<T> {
375    receiver: PriorityQueueReceiver<T>,
376    ended: bool,
377}
378impl<T> Iterator for TryIter<T> {
379    type Item = Result<T, RecvError>;
380
381    fn next(&mut self) -> Option<Self::Item> {
382        if self.ended {
383            return None;
384        }
385
386        let res = self.receiver.try_pop();
387        self.ended = res.is_err();
388
389        res.transpose()
390    }
391}
392impl<T> FusedIterator for TryIter<T> {}
393
394#[cfg(test)]
395mod tests {
396    use collections::HashSet;
397
398    use super::*;
399
400    #[test]
401    fn all_tasks_get_yielded() {
402        let (tx, mut rx) = PriorityQueueReceiver::new();
403        tx.send(Priority::Medium, 20).unwrap();
404        tx.send(Priority::High, 30).unwrap();
405        tx.send(Priority::Low, 10).unwrap();
406        tx.send(Priority::Medium, 21).unwrap();
407        tx.send(Priority::High, 31).unwrap();
408
409        drop(tx);
410
411        assert_eq!(
412            rx.iter().collect::<HashSet<_>>(),
413            [30, 31, 20, 21, 10].into_iter().collect::<HashSet<_>>()
414        )
415    }
416
417    #[test]
418    fn new_high_prio_task_get_scheduled_quickly() {
419        let (tx, mut rx) = PriorityQueueReceiver::new();
420        for _ in 0..100 {
421            tx.send(Priority::Low, 1).unwrap();
422        }
423
424        assert_eq!(rx.pop().unwrap(), 1);
425        tx.send(Priority::High, 3).unwrap();
426        assert_eq!(rx.pop().unwrap(), 3);
427        assert_eq!(rx.pop().unwrap(), 1);
428    }
429}
430