dast.async.thread source code

1 module dast.async.thread;
2 
3 import core.atomic,
4 core.thread;
5 
6 private struct TaskBase {
7 	void function(void*) run;
8 }
9 
10 struct Task(alias fun, Args...) {
11 private:
12 	TaskBase base = TaskBase(&impl);
13 	Args _args;
14 
15 	static if (Args.length) {
16 		this(Args args) {
17 			_args = args;
18 		}
19 	}
20 
21 	static void impl(void* p) {
22 		fun((cast(typeof(this)*)p)._args);
23 	}
24 }
25 
26 final class TaskThread : Thread {
27 	this(void delegate() dg) {
28 		super(dg);
29 	}
30 
31 	void run() nothrow {
32 		running.atomicStore(true);
33 		super.start();
34 	}
35 
36 	shared bool running;
37 }
38 
39 struct ThreadPool {
40 	import core.memory,
41 	core.sync,
42 	tame.fixed.queue;
43 
44 	enum State : ubyte {
45 		running,
46 		done,
47 		stop
48 	}
49 
50 	private {
51 		TaskThread[] pool;
52 		Mutex mutex;
53 		Condition cond;
54 		shared State status;
55 		Queue!(TaskBase*, 2048 / size_t.sizeof) queue;
56 		shared uint nWorkers;
57 	}
58 	immutable uint timeoutMs;
59 
60 	this(uint size, uint timeout_ms = 0) @trusted {
61 		timeoutMs = timeout_ms;
62 		pool = (cast(TaskThread*)pureMalloc(size * TaskThread.sizeof))[0 .. size];
63 		// TODO: check pool is null
64 		nWorkers = size;
65 		mutex = new Mutex();
66 		cond = new Condition(mutex);
67 		foreach (ref t; pool) {
68 			t = new TaskThread(&workLoop);
69 			t.run();
70 		}
71 	}
72 
73 	~this() @trusted {
74 		pureFree(pool.ptr);
75 		finish(true);
76 	}
77 
78 	@property size() const => pool.length;
79 
80 	@property state() const @trusted => atomicLoad(status);
81 
82 	void workLoop() {
83 		loop: while (atomicLoad(status) != State.stop) {
84 			mutex.lock_nothrow();
85 			while (queue.empty) {
86 				if (timeoutMs) {
87 					if (!cond.wait(timeoutMs.msecs)) {
88 						mutex.unlock_nothrow();
89 						break loop;
90 					}
91 				} else
92 					cond.wait();
93 			}
94 			auto task = queue.pop();
95 			cond.notify();
96 			mutex.unlock_nothrow();
97 			task.run(task);
98 		}
99 		auto tthis = Thread.getThis();
100 		foreach (t; pool) {
101 			if (t == tthis) {
102 				t.running.atomicStore(false);
103 				atomicStore(nWorkers, nWorkers - 1);
104 				break;
105 			}
106 		}
107 	}
108 
109 	void finish(bool blocking = false) @trusted {
110 		{
111 			mutex.lock_nothrow();
112 			scope (exit)
113 				mutex.unlock_nothrow();
114 			cas(&status, State.running, State.done);
115 			cond.notifyAll();
116 		}
117 		if (blocking) {
118 			foreach (t; pool) {
119 				t.join();
120 			}
121 		}
122 	}
123 
124 	void stop() @trusted {
125 		mutex.lock_nothrow();
126 		scope (exit)
127 			mutex.unlock_nothrow();
128 		cas(&status, State.running, State.stop);
129 		cond.notifyAll();
130 	}
131 
132 	void run(alias fn, Args...)(auto ref Args args) {
133 		run(&new Task!(fn, Args)(args).base);
134 	}
135 
136 	private void run(TaskBase* task)
137 	in (task) {
138 		mutex.lock_nothrow();
139 		scope (exit)
140 			mutex.unlock_nothrow();
141 		if (queue.full)
142 			cond.wait();
143 		if (atomicLoad(status) != State.running)
144 			return;
145 		queue.push(task);
146 		if (timeoutMs && atomicLoad(nWorkers) < pool.length) {
147 			foreach (t; pool) {
148 				if (!atomicLoad(t.running)) {
149 					t.run();
150 					atomicStore(nWorkers, nWorkers + 1);
151 					break;
152 				}
153 			}
154 		}
155 		cond.notify();
156 	}
157 }