-
Notifications
You must be signed in to change notification settings - Fork 0
/
thread_pool.py
92 lines (77 loc) · 2.88 KB
/
thread_pool.py
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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Author: Yuande <miraclecome (at) gmail.com>
# This code is under Creative Commons CC BY-NC-SA license
# http://creativecommons.org/licenses/by-nc-sa/3.0/
import queue
import threading
import log
class worker(threading.Thread):
worker_count = 0
def __init__(self, work_queue, result_queue, timeout=0, **kwargs):
threading.Thread.__init__(self, **kwargs)
self.id = worker.worker_count
worker.worker_count += 1
self.daemon = True
self.work_queue = work_queue
self.result_queue = result_queue
self.timeout = timeout
def run(self):
''' the get-and-do loop'''
while True:
try:
callable, args, kwargs = self.work_queue.get(timeout=self.timeout)
result = callable(*args, **kwargs)
if result: self.result_queue.put(result) # only add url_graph result, filter indexing result
# print('worker[%d]: %s' % (self.id, str(result)))
# formally need task_done(), but we can omit it here.
self.work_queue.task_done()
except queue.Empty:
break
except:
message = 'worker[{0}]'.format(self.id)
log.log_traceback(msg=message)
class thread_pool(object):
def __init__(self, num_workers=10, timeout=10):
self.work_queue = queue.Queue()
self.result_queue = queue.Queue()
self.workers = []
self.timeout = timeout
self._recruitThreads(num_workers)
def _recruitThreads(self, num_workers):
for i in range(num_workers):
worker_obj = worker(self.work_queue, self.result_queue, self.timeout)
self.workers.append(worker_obj)
def start(self):
for w in self.workers:
w.start()
def wait_completion(self):
''' waiting for each worker to terminate. '''
while self.workers:
worker_obj = self.workers.pop()
worker_obj.join() # blocking, waiting for the all complete.
if worker_obj.is_alive() and not self.work_queue.Empty:
self.workers.append(worker_obj)
def add_job(self, callable, *args, **kwargs):
self.work_queue.put( (callable, args, kwargs) )
def get_one_result(self, *args, **kwargs):
if not self.result_queue.empty():
return self.result_queue.get(*args, **kwargs)
return None
def ssize(self):
return len(self.workers)
if __name__ == '__main__':
import md5
obj = thread_pool(3)
for i in range(10):
obj.add_job(md5.md5_str, 'who are you')
obj.start()
print(obj.ssize())
for i in range(10):
obj.add_job(md5.md5_str, 'where are we')
obj.wait_completion()
ret = obj.get_one_result()
while ret:
print(ret)
ret = obj.get_one_result()
print(obj.ssize())