1use std::sync::{Arc, RwLock, RwLockWriteGuard};
2use std::sync::PoisonError;
3use std::sync::mpsc;
4use std::sync::mpsc::{Sender};
5use std::thread::JoinHandle;
6use std::collections::HashMap;
7use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
8
9#[derive(Clone, Copy, Eq, PartialOrd, PartialEq, Debug)]
10pub enum Notify {
11 Started(usize),
12 Go,
13 Terminated(usize),
14 Quit,
15}
16pub struct ThreadQueue {
17 sender_map:Arc<RwLock<HashMap<usize,(Sender<Notify>,bool)>>>,
18 sender:Sender<Notify>,
19 last_id:Arc<RwLock<usize>>,
20 working:Arc<AtomicBool>,
21 busy_count:Arc<AtomicUsize>,
22 worker_handler:Option<JoinHandle<()>>
23}
24impl ThreadQueue {
25 pub fn new() -> ThreadQueue {
26 let (ss,sr) = mpsc::channel();
27
28 let last_id = Arc::new(RwLock::new(0));
29 let working = Arc::new(AtomicBool::new(true));
30 let busy_count = Arc::new(AtomicUsize::new(0));
31
32 let sender_map = Arc::new(RwLock::new(HashMap::<usize,(Sender<Notify>,bool)>::new()));
33
34 let h = {
35 let sender_map = Arc::clone(&sender_map);
36 let working = Arc::clone(&working);
37 let busy_count = Arc::clone(&busy_count);
38
39 let mut current_id = 0usize;
40
41 let (ws,wr) = mpsc::channel();
42
43 let h = std::thread::spawn(move || {
44 ws.send(()).unwrap();
45
46 loop {
47 match sr.recv().unwrap() {
48 Notify::Started(id) => {
49 if id == current_id {
50 match sender_map.read().unwrap().get(&id) {
51 Some((ref sender, false)) => {
52 sender.send(Notify::Go).unwrap();
53 },
54 _ => ()
55 }
56
57 }
58 },
59 Notify::Terminated(id) => {
60 if id == current_id {
61 match sender_map.write() {
62 Ok(mut map) => {
63 map.remove(&id);
64 let id = id + 1;
65 current_id = id;
66 match map.get_mut(&id) {
67 Some((ref sender, ref mut started)) => {
68 *started = true;
69 sender.send(Notify::Go).unwrap();
70 },
71 None => ()
72 }
73 },
74 Err(ref e) => {
75 panic!("{:?}", e);
76 }
77 };
78
79 busy_count.fetch_sub(1, Ordering::Release);
80
81 if !working.load(Ordering::Acquire) && busy_count.load(Ordering::Acquire) == 0 {
82 break;
83 }
84 }
85 },
86 Notify::Quit if busy_count.load(Ordering::Acquire) > 0 => {
87
88 },
89 Notify::Quit if sender_map.read().unwrap().is_empty() => {
90 break;
91 },
92 Notify::Quit => {
93 panic!("There are still threads waiting to be processed.");
94 }
95 _ => (),
96 }
97 }
98 });
99
100 wr.recv().unwrap();
101
102 h
103 };
104
105 ThreadQueue {
106 sender_map:sender_map,
107 sender:ss,
108 last_id:last_id,
109 working:working,
110 busy_count:Arc::clone(&busy_count),
111 worker_handler:Some(h)
112 }
113 }
114
115 pub fn submit<F,T>(&mut self,f:F) ->
116 Result<JoinHandle<T>,PoisonError<RwLockWriteGuard<'_,HashMap<usize,(Sender<Notify>,bool)>>>>
117 where F: FnOnce() -> T, F: Send + 'static, T: Send + 'static {
118
119 let (cs,cr) = mpsc::channel();
120
121 match self.last_id.write().unwrap() {
122 mut last_id => {
123 {
124 self.sender_map.write()?.insert(*last_id, (cs.clone(),false));
125 }
126
127 let ss = self.sender.clone();
128
129 let id = *last_id;
130
131 let (on_start_sender,on_start_receiver) = mpsc::channel();
132
133 let busy_count = Arc::clone(&self.busy_count);
134
135 let r = std::thread::spawn(move || {
136 ss.send(Notify::Started(id)).unwrap();
137 busy_count.fetch_add(1, Ordering::Release);
138 on_start_sender.send(()).unwrap();
139 cr.recv().unwrap();
140
141 let r = f();
142
143 ss.send(Notify::Terminated(id)).unwrap();
144 r
145 });
146
147 *last_id += 1;
148
149 on_start_receiver.recv().unwrap();
150
151 Ok(r)
152 }
153 }
154 }
155}
156impl Drop for ThreadQueue {
157 fn drop(&mut self) {
158 let r = self.sender.send(Notify::Quit);
159 self.working.store(false, Ordering::Release);
160 self.worker_handler.take().map(|h| {
161 h.join().unwrap()
162 }).unwrap();
163 r.unwrap();
164 }
165}