blob: 5d0165c760d42b1fd885e24577b32b26054ecb8e (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
|
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::sync_channel;
use std::sync::mpsc::{Receiver, SyncSender};
use spin::Mutex;
pub struct ParallelQueue<T> {
next: AtomicUsize, // next round-robin index
queues: Vec<Mutex<SyncSender<T>>>, // work queues (1 per thread)
}
impl<T> ParallelQueue<T> {
pub fn new(queues: usize) -> (Vec<Receiver<T>>, Self) {
let mut rxs = vec![];
let mut txs = vec![];
for _ in 0..queues {
let (tx, rx) = sync_channel(128);
txs.push(Mutex::new(tx));
rxs.push(rx);
}
(
rxs,
ParallelQueue {
next: AtomicUsize::new(0),
queues: txs,
},
)
}
pub fn send(&self, v: T) {
let len = self.queues.len();
let idx = self.next.fetch_add(1, Ordering::SeqCst);
let que = self.queues[idx % len].lock();
que.send(v).unwrap();
}
pub fn close(&self) {
for i in 0..self.queues.len() {
let (tx, _) = sync_channel(0);
let queue = &self.queues[i];
*queue.lock() = tx;
}
}
}
|