Skip to main content

queuingtask/
lib.rs

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}