gpui/
queue.rs

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    /// Returns the number of queued elements across all priorities.
229    pub(crate) fn len(&self) -> usize {
230        let queues = self.state.queues.lock();
231        queues.high_priority.len() + queues.medium_priority.len() + queues.low_priority.len()
232    }
233
234    /// Tries to pop one element from the priority queue without blocking.
235    ///
236    /// This will early return if there are no elements in the queue.
237    ///
238    /// This method is best suited if you only intend to pop one element, for better performance
239    /// on large queues see [`Self::try_iter`]
240    ///
241    /// # Errors
242    ///
243    /// If the sender was dropped
244    pub fn try_pop(&mut self) -> Result<Option<T>, RecvError> {
245        self.pop_inner(false)
246    }
247
248    pub fn spin_try_pop(&mut self) -> Result<Option<T>, RecvError> {
249        use Priority as P;
250
251        let Some(mut queues) = self.state.spin_try_recv()? else {
252            return Ok(None);
253        };
254
255        let high = P::High.weight() * !queues.high_priority.is_empty() as u32;
256        let medium = P::Medium.weight() * !queues.medium_priority.is_empty() as u32;
257        let low = P::Low.weight() * !queues.low_priority.is_empty() as u32;
258        let mut mass = high + medium + low;
259
260        if !queues.high_priority.is_empty() {
261            let flip = self.rand.random_ratio(P::High.weight(), mass);
262            if flip {
263                return Ok(queues.high_priority.pop_front());
264            }
265            mass -= P::High.weight();
266        }
267
268        if !queues.medium_priority.is_empty() {
269            let flip = self.rand.random_ratio(P::Medium.weight(), mass);
270            if flip {
271                return Ok(queues.medium_priority.pop_front());
272            }
273            mass -= P::Medium.weight();
274        }
275
276        if !queues.low_priority.is_empty() {
277            let flip = self.rand.random_ratio(P::Low.weight(), mass);
278            if flip {
279                return Ok(queues.low_priority.pop_front());
280            }
281        }
282
283        Ok(None)
284    }
285
286    /// Pops an element from the priority queue blocking if necessary.
287    ///
288    /// This method is best suited if you only intend to pop one element, for better performance
289    /// on large queues see [`Self::iter``]
290    ///
291    /// # Errors
292    ///
293    /// If the sender was dropped
294    pub fn pop(&mut self) -> Result<T, RecvError> {
295        self.pop_inner(true).map(|e| e.unwrap())
296    }
297
298    /// Returns an iterator over the elements of the queue
299    /// this iterator will end when all elements have been consumed and will not wait for new ones.
300    pub fn try_iter(self) -> TryIter<T> {
301        TryIter {
302            receiver: self,
303            ended: false,
304        }
305    }
306
307    /// Returns an iterator over the elements of the queue
308    /// this iterator will wait for new elements if the queue is empty.
309    pub fn iter(self) -> Iter<T> {
310        Iter(self)
311    }
312
313    #[inline(always)]
314    // algorithm is the loaded die from biased coin from
315    // https://www.keithschwarz.com/darts-dice-coins/
316    fn pop_inner(&mut self, block: bool) -> Result<Option<T>, RecvError> {
317        use Priority as P;
318
319        let mut queues = if !block {
320            let Some(queues) = self.state.try_recv()? else {
321                return Ok(None);
322            };
323            queues
324        } else {
325            self.state.recv()?
326        };
327
328        let high = P::High.weight() * !queues.high_priority.is_empty() as u32;
329        let medium = P::Medium.weight() * !queues.medium_priority.is_empty() as u32;
330        let low = P::Low.weight() * !queues.low_priority.is_empty() as u32;
331        let mut mass = high + medium + low; //%
332
333        if !queues.high_priority.is_empty() {
334            let flip = self.rand.random_ratio(P::High.weight(), mass);
335            if flip {
336                return Ok(queues.high_priority.pop_front());
337            }
338            mass -= P::High.weight();
339        }
340
341        if !queues.medium_priority.is_empty() {
342            let flip = self.rand.random_ratio(P::Medium.weight(), mass);
343            if flip {
344                return Ok(queues.medium_priority.pop_front());
345            }
346            mass -= P::Medium.weight();
347        }
348
349        if !queues.low_priority.is_empty() {
350            let flip = self.rand.random_ratio(P::Low.weight(), mass);
351            if flip {
352                return Ok(queues.low_priority.pop_front());
353            }
354        }
355
356        Ok(None)
357    }
358}
359
360impl<T> Drop for PriorityQueueReceiver<T> {
361    fn drop(&mut self) {
362        self.state
363            .receiver_count
364            .fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
365    }
366}
367
368#[doc(hidden)]
369pub struct Iter<T>(PriorityQueueReceiver<T>);
370impl<T> Iterator for Iter<T> {
371    type Item = T;
372
373    fn next(&mut self) -> Option<Self::Item> {
374        self.0.pop().ok()
375    }
376}
377impl<T> FusedIterator for Iter<T> {}
378
379#[doc(hidden)]
380pub struct TryIter<T> {
381    receiver: PriorityQueueReceiver<T>,
382    ended: bool,
383}
384impl<T> Iterator for TryIter<T> {
385    type Item = Result<T, RecvError>;
386
387    fn next(&mut self) -> Option<Self::Item> {
388        if self.ended {
389            return None;
390        }
391
392        let res = self.receiver.try_pop();
393        self.ended = res.is_err();
394
395        res.transpose()
396    }
397}
398impl<T> FusedIterator for TryIter<T> {}
399
400#[cfg(test)]
401mod tests {
402    use collections::HashSet;
403
404    use super::*;
405
406    #[test]
407    fn all_tasks_get_yielded() {
408        let (tx, mut rx) = PriorityQueueReceiver::new();
409        tx.send(Priority::Medium, 20).unwrap();
410        tx.send(Priority::High, 30).unwrap();
411        tx.send(Priority::Low, 10).unwrap();
412        tx.send(Priority::Medium, 21).unwrap();
413        tx.send(Priority::High, 31).unwrap();
414
415        drop(tx);
416
417        assert_eq!(
418            rx.iter().collect::<HashSet<_>>(),
419            [30, 31, 20, 21, 10].into_iter().collect::<HashSet<_>>()
420        )
421    }
422
423    #[test]
424    fn new_high_prio_task_get_scheduled_quickly() {
425        let (tx, mut rx) = PriorityQueueReceiver::new();
426        for _ in 0..100 {
427            tx.send(Priority::Low, 1).unwrap();
428        }
429
430        assert_eq!(rx.pop().unwrap(), 1);
431        tx.send(Priority::High, 3).unwrap();
432        assert_eq!(rx.pop().unwrap(), 3);
433        assert_eq!(rx.pop().unwrap(), 1);
434    }
435}
gpui/queue.rs

gpui/
queue.rs
