gpui/
queue.rs

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

gpui/
queue.rs
