queues.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395
  1. #
  2. # Module implementing queues
  3. #
  4. # multiprocessing/queues.py
  5. #
  6. # Copyright (c) 2006-2008, R Oudkerk
  7. # All rights reserved.
  8. #
  9. # Redistribution and use in source and binary forms, with or without
  10. # modification, are permitted provided that the following conditions
  11. # are met:
  12. #
  13. # 1. Redistributions of source code must retain the above copyright
  14. # notice, this list of conditions and the following disclaimer.
  15. # 2. Redistributions in binary form must reproduce the above copyright
  16. # notice, this list of conditions and the following disclaimer in the
  17. # documentation and/or other materials provided with the distribution.
  18. # 3. Neither the name of author nor the names of any contributors may be
  19. # used to endorse or promote products derived from this software
  20. # without specific prior written permission.
  21. #
  22. # THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS "AS IS" AND
  23. # ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  24. # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  25. # ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
  26. # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  27. # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
  28. # OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
  29. # HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  30. # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
  31. # OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
  32. # SUCH DAMAGE.
  33. #
  34. __all__ = ['Queue', 'SimpleQueue', 'JoinableQueue']
  35. import sys
  36. import os
  37. import threading
  38. import collections
  39. import time
  40. import atexit
  41. import weakref
  42. from Queue import Empty, Full
  43. import _multiprocessing
  44. from . import Pipe
  45. from .synchronize import Lock, BoundedSemaphore, Semaphore, Condition
  46. from .util import debug, info, Finalize, register_after_fork, is_exiting
  47. from .forking import assert_spawning
  48. #
  49. # Queue type using a pipe, buffer and thread
  50. #
  51. class Queue(object):
  52. def __init__(self, maxsize=0):
  53. if maxsize <= 0:
  54. maxsize = _multiprocessing.SemLock.SEM_VALUE_MAX
  55. self._maxsize = maxsize
  56. self._reader, self._writer = Pipe(duplex=False)
  57. self._rlock = Lock()
  58. self._opid = os.getpid()
  59. if sys.platform == 'win32':
  60. self._wlock = None
  61. else:
  62. self._wlock = Lock()
  63. self._sem = BoundedSemaphore(maxsize)
  64. self._after_fork()
  65. if sys.platform != 'win32':
  66. register_after_fork(self, Queue._after_fork)
  67. def __getstate__(self):
  68. assert_spawning(self)
  69. return (self._maxsize, self._reader, self._writer,
  70. self._rlock, self._wlock, self._sem, self._opid)
  71. def __setstate__(self, state):
  72. (self._maxsize, self._reader, self._writer,
  73. self._rlock, self._wlock, self._sem, self._opid) = state
  74. self._after_fork()
  75. def _after_fork(self):
  76. debug('Queue._after_fork()')
  77. self._notempty = threading.Condition(threading.Lock())
  78. self._buffer = collections.deque()
  79. self._thread = None
  80. self._jointhread = None
  81. self._joincancelled = False
  82. self._closed = False
  83. self._close = None
  84. self._send = self._writer.send
  85. self._recv = self._reader.recv
  86. self._poll = self._reader.poll
  87. def put(self, obj, block=True, timeout=None):
  88. assert not self._closed
  89. if not self._sem.acquire(block, timeout):
  90. raise Full
  91. self._notempty.acquire()
  92. try:
  93. if self._thread is None:
  94. self._start_thread()
  95. self._buffer.append(obj)
  96. self._notempty.notify()
  97. finally:
  98. self._notempty.release()
  99. def get(self, block=True, timeout=None):
  100. if block and timeout is None:
  101. self._rlock.acquire()
  102. try:
  103. res = self._recv()
  104. self._sem.release()
  105. return res
  106. finally:
  107. self._rlock.release()
  108. else:
  109. if block:
  110. deadline = time.time() + timeout
  111. if not self._rlock.acquire(block, timeout):
  112. raise Empty
  113. try:
  114. if block:
  115. timeout = deadline - time.time()
  116. if timeout < 0 or not self._poll(timeout):
  117. raise Empty
  118. elif not self._poll():
  119. raise Empty
  120. res = self._recv()
  121. self._sem.release()
  122. return res
  123. finally:
  124. self._rlock.release()
  125. def qsize(self):
  126. # Raises NotImplementedError on Mac OSX because of broken sem_getvalue()
  127. return self._maxsize - self._sem._semlock._get_value()
  128. def empty(self):
  129. return not self._poll()
  130. def full(self):
  131. return self._sem._semlock._is_zero()
  132. def get_nowait(self):
  133. return self.get(False)
  134. def put_nowait(self, obj):
  135. return self.put(obj, False)
  136. def close(self):
  137. self._closed = True
  138. try:
  139. self._reader.close()
  140. finally:
  141. close = self._close
  142. if close:
  143. self._close = None
  144. close()
  145. def join_thread(self):
  146. debug('Queue.join_thread()')
  147. assert self._closed
  148. if self._jointhread:
  149. self._jointhread()
  150. def cancel_join_thread(self):
  151. debug('Queue.cancel_join_thread()')
  152. self._joincancelled = True
  153. try:
  154. self._jointhread.cancel()
  155. except AttributeError:
  156. pass
  157. def _start_thread(self):
  158. debug('Queue._start_thread()')
  159. # Start thread which transfers data from buffer to pipe
  160. self._buffer.clear()
  161. self._thread = threading.Thread(
  162. target=Queue._feed,
  163. args=(self._buffer, self._notempty, self._send,
  164. self._wlock, self._writer.close),
  165. name='QueueFeederThread'
  166. )
  167. self._thread.daemon = True
  168. debug('doing self._thread.start()')
  169. self._thread.start()
  170. debug('... done self._thread.start()')
  171. # On process exit we will wait for data to be flushed to pipe.
  172. if not self._joincancelled:
  173. self._jointhread = Finalize(
  174. self._thread, Queue._finalize_join,
  175. [weakref.ref(self._thread)],
  176. exitpriority=-5
  177. )
  178. # Send sentinel to the thread queue object when garbage collected
  179. self._close = Finalize(
  180. self, Queue._finalize_close,
  181. [self._buffer, self._notempty],
  182. exitpriority=10
  183. )
  184. @staticmethod
  185. def _finalize_join(twr):
  186. debug('joining queue thread')
  187. thread = twr()
  188. if thread is not None:
  189. thread.join()
  190. debug('... queue thread joined')
  191. else:
  192. debug('... queue thread already dead')
  193. @staticmethod
  194. def _finalize_close(buffer, notempty):
  195. debug('telling queue thread to quit')
  196. notempty.acquire()
  197. try:
  198. buffer.append(_sentinel)
  199. notempty.notify()
  200. finally:
  201. notempty.release()
  202. @staticmethod
  203. def _feed(buffer, notempty, send, writelock, close):
  204. debug('starting thread to feed data to pipe')
  205. nacquire = notempty.acquire
  206. nrelease = notempty.release
  207. nwait = notempty.wait
  208. bpopleft = buffer.popleft
  209. sentinel = _sentinel
  210. if sys.platform != 'win32':
  211. wacquire = writelock.acquire
  212. wrelease = writelock.release
  213. else:
  214. wacquire = None
  215. try:
  216. while 1:
  217. nacquire()
  218. try:
  219. if not buffer:
  220. nwait()
  221. finally:
  222. nrelease()
  223. try:
  224. while 1:
  225. obj = bpopleft()
  226. if obj is sentinel:
  227. debug('feeder thread got sentinel -- exiting')
  228. close()
  229. return
  230. if wacquire is None:
  231. send(obj)
  232. else:
  233. wacquire()
  234. try:
  235. send(obj)
  236. finally:
  237. wrelease()
  238. except IndexError:
  239. pass
  240. except Exception, e:
  241. # Since this runs in a daemon thread the resources it uses
  242. # may be become unusable while the process is cleaning up.
  243. # We ignore errors which happen after the process has
  244. # started to cleanup.
  245. try:
  246. if is_exiting():
  247. info('error in queue thread: %s', e)
  248. else:
  249. import traceback
  250. traceback.print_exc()
  251. except Exception:
  252. pass
  253. _sentinel = object()
  254. #
  255. # A queue type which also supports join() and task_done() methods
  256. #
  257. # Note that if you do not call task_done() for each finished task then
  258. # eventually the counter's semaphore may overflow causing Bad Things
  259. # to happen.
  260. #
  261. class JoinableQueue(Queue):
  262. def __init__(self, maxsize=0):
  263. Queue.__init__(self, maxsize)
  264. self._unfinished_tasks = Semaphore(0)
  265. self._cond = Condition()
  266. def __getstate__(self):
  267. return Queue.__getstate__(self) + (self._cond, self._unfinished_tasks)
  268. def __setstate__(self, state):
  269. Queue.__setstate__(self, state[:-2])
  270. self._cond, self._unfinished_tasks = state[-2:]
  271. def put(self, obj, block=True, timeout=None):
  272. assert not self._closed
  273. if not self._sem.acquire(block, timeout):
  274. raise Full
  275. self._notempty.acquire()
  276. self._cond.acquire()
  277. try:
  278. if self._thread is None:
  279. self._start_thread()
  280. self._buffer.append(obj)
  281. self._unfinished_tasks.release()
  282. self._notempty.notify()
  283. finally:
  284. self._cond.release()
  285. self._notempty.release()
  286. def task_done(self):
  287. self._cond.acquire()
  288. try:
  289. if not self._unfinished_tasks.acquire(False):
  290. raise ValueError('task_done() called too many times')
  291. if self._unfinished_tasks._semlock._is_zero():
  292. self._cond.notify_all()
  293. finally:
  294. self._cond.release()
  295. def join(self):
  296. self._cond.acquire()
  297. try:
  298. if not self._unfinished_tasks._semlock._is_zero():
  299. self._cond.wait()
  300. finally:
  301. self._cond.release()
  302. #
  303. # Simplified Queue type -- really just a locked pipe
  304. #
  305. class SimpleQueue(object):
  306. def __init__(self):
  307. self._reader, self._writer = Pipe(duplex=False)
  308. self._rlock = Lock()
  309. if sys.platform == 'win32':
  310. self._wlock = None
  311. else:
  312. self._wlock = Lock()
  313. self._make_methods()
  314. def empty(self):
  315. return not self._reader.poll()
  316. def __getstate__(self):
  317. assert_spawning(self)
  318. return (self._reader, self._writer, self._rlock, self._wlock)
  319. def __setstate__(self, state):
  320. (self._reader, self._writer, self._rlock, self._wlock) = state
  321. self._make_methods()
  322. def _make_methods(self):
  323. recv = self._reader.recv
  324. racquire, rrelease = self._rlock.acquire, self._rlock.release
  325. def get():
  326. racquire()
  327. try:
  328. return recv()
  329. finally:
  330. rrelease()
  331. self.get = get
  332. if self._wlock is None:
  333. # writes to a message oriented win32 pipe are atomic
  334. self.put = self._writer.send
  335. else:
  336. send = self._writer.send
  337. wacquire, wrelease = self._wlock.acquire, self._wlock.release
  338. def put(obj):
  339. wacquire()
  340. try:
  341. return send(obj)
  342. finally:
  343. wrelease()
  344. self.put = put