komora_sync/
segmented_priority_queue.rs

1use std::collections::BinaryHeap;
2use std::sync::{Condvar, Mutex};
3
4use crate::Prioritized;
5
6/// Like a [`PriorityQueue`] but avoids starvation by
7/// using a rotating bip-buffer of priority queues.
8#[derive(Debug)]
9pub struct SegmentedPriorityQueue<T> {
10    reader: Mutex<BinaryHeap<Prioritized<T>>>,
11    writer: Mutex<BinaryHeap<Prioritized<T>>>,
12    cv: Condvar,
13}
14
15impl<T> Default for SegmentedPriorityQueue<T> {
16    fn default() -> Self {
17        Self::new()
18    }
19}
20
21impl<T> SegmentedPriorityQueue<T> {
22    pub fn new() -> SegmentedPriorityQueue<T> {
23        SegmentedPriorityQueue {
24            reader: Mutex::default(),
25            writer: Mutex::default(),
26            cv: Condvar::new(),
27        }
28    }
29
30    /// Higher priority tends to gets popped first.
31    ///
32    /// The internal bip buffer of priority queues gets
33    /// rotated when the read side is empty.
34    ///
35    /// # Examples
36    /// ```
37    /// let pq = komora_sync::SegmentedPriorityQueue::new();
38    /// pq.push(2, 2);
39    /// pq.push(1, 1);
40    ///
41    /// // internal state:
42    /// // read buffer: []
43    /// // write buffer: [2, 1]
44    ///
45    /// // queue rotated on pop when read buffer is empty
46    /// assert_eq!(pq.pop(), 2);
47    ///
48    /// // internal state:
49    /// // read buffer: [1]
50    /// // write buffer: []
51    ///
52    /// // new writes get pushed to write buffer
53    /// pq.push(4, 4);
54    /// pq.push(3, 3);
55    ///
56    /// // internal state:
57    /// // read buffer: [1]
58    /// // write buffer: [4, 3]
59    ///
60    /// // this is the last element in the reader side
61    /// assert_eq!(pq.pop(), 1);
62    ///
63    /// // queue rotated on pop when reader is empty
64    /// assert_eq!(pq.pop(), 4);
65    /// assert_eq!(pq.pop(), 3);
66    /// ```
67    pub fn push(&self, t: T, priority: u64) {
68        let mut q = self.writer.lock().unwrap();
69        q.push(Prioritized { t, priority });
70        drop(q);
71        self.cv.notify_one();
72    }
73
74    pub fn pop(&self) -> T {
75        let mut q = self.reader.lock().unwrap();
76
77        while q.is_empty() {
78            // rotate the bip buffer
79            let mut q2 = self.writer.lock().unwrap();
80
81            std::mem::swap(&mut *q, &mut *q2);
82
83            drop(q2);
84
85            if q.is_empty() {
86                q = self.cv.wait(q).unwrap();
87            }
88        }
89
90        q.pop().unwrap().t
91    }
92}
93
94#[test]
95fn segmented_priority_queue_concurrent_pops() {
96    use std::sync::atomic::{AtomicU64, Ordering};
97    use std::sync::{Arc, Barrier};
98
99    const N_WRITES_PER_WRITER: u64 = 1_000_000;
100    const N_WRITERS: u64 = 10;
101    const N_WRITES: u64 = N_WRITERS * N_WRITES_PER_WRITER;
102    const N_READERS: u64 = 10;
103
104    let q = Arc::new(SegmentedPriorityQueue::new());
105    let barrier = Arc::new(Barrier::new(N_WRITERS.try_into().unwrap()));
106    let read = Arc::new(AtomicU64::new(0));
107
108    for i in 0..N_WRITERS {
109        let q = q.clone();
110        let barrier = barrier.clone();
111        std::thread::spawn(move || {
112            barrier.wait();
113
114            for j in 0..N_WRITES_PER_WRITER {
115                q.push(i * j, i);
116            }
117        });
118    }
119
120    for _ in 0..N_READERS {
121        let q = q.clone();
122        let read = read.clone();
123        std::thread::spawn(move || loop {
124            q.pop();
125            read.fetch_add(1, Ordering::SeqCst);
126        });
127    }
128
129    while read.load(Ordering::SeqCst) < N_WRITES {
130        std::hint::spin_loop();
131    }
132}
komora_sync/segmented_priority_queue.rs

komora_sync/
segmented_priority_queue.rs
