queues.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390
  1. #
  2. # Module implementing queues
  3. #
  4. # multiprocessing/queues.py
  5. #
  6. # Copyright (c) 2006-2008, R Oudkerk
  7. # Licensed to PSF under a Contributor Agreement.
  8. #
  9. from __future__ import absolute_import
  10. import sys
  11. import os
  12. import threading
  13. import collections
  14. import weakref
  15. import errno
  16. from . import connection
  17. from . import context
  18. from .compat import get_errno
  19. from .five import monotonic, Empty, Full
  20. from .util import (
  21. debug, error, info, Finalize, register_after_fork, is_exiting,
  22. )
  23. from .reduction import ForkingPickler
  24. __all__ = ['Queue', 'SimpleQueue', 'JoinableQueue']
  25. class Queue(object):
  26. '''
  27. Queue type using a pipe, buffer and thread
  28. '''
  29. def __init__(self, maxsize=0, *args, **kwargs):
  30. try:
  31. ctx = kwargs['ctx']
  32. except KeyError:
  33. raise TypeError('missing 1 required keyword-only argument: ctx')
  34. if maxsize <= 0:
  35. # Can raise ImportError (see issues #3770 and #23400)
  36. from .synchronize import SEM_VALUE_MAX as maxsize # noqa
  37. self._maxsize = maxsize
  38. self._reader, self._writer = connection.Pipe(duplex=False)
  39. self._rlock = ctx.Lock()
  40. self._opid = os.getpid()
  41. if sys.platform == 'win32':
  42. self._wlock = None
  43. else:
  44. self._wlock = ctx.Lock()
  45. self._sem = ctx.BoundedSemaphore(maxsize)
  46. # For use by concurrent.futures
  47. self._ignore_epipe = False
  48. self._after_fork()
  49. if sys.platform != 'win32':
  50. register_after_fork(self, Queue._after_fork)
  51. def __getstate__(self):
  52. context.assert_spawning(self)
  53. return (self._ignore_epipe, self._maxsize, self._reader, self._writer,
  54. self._rlock, self._wlock, self._sem, self._opid)
  55. def __setstate__(self, state):
  56. (self._ignore_epipe, self._maxsize, self._reader, self._writer,
  57. self._rlock, self._wlock, self._sem, self._opid) = state
  58. self._after_fork()
  59. def _after_fork(self):
  60. debug('Queue._after_fork()')
  61. self._notempty = threading.Condition(threading.Lock())
  62. self._buffer = collections.deque()
  63. self._thread = None
  64. self._jointhread = None
  65. self._joincancelled = False
  66. self._closed = False
  67. self._close = None
  68. self._send_bytes = self._writer.send
  69. self._recv = self._reader.recv
  70. self._send_bytes = self._writer.send_bytes
  71. self._recv_bytes = self._reader.recv_bytes
  72. self._poll = self._reader.poll
  73. def put(self, obj, block=True, timeout=None):
  74. assert not self._closed
  75. if not self._sem.acquire(block, timeout):
  76. raise Full
  77. with self._notempty:
  78. if self._thread is None:
  79. self._start_thread()
  80. self._buffer.append(obj)
  81. self._notempty.notify()
  82. def get(self, block=True, timeout=None):
  83. if block and timeout is None:
  84. with self._rlock:
  85. res = self._recv_bytes()
  86. self._sem.release()
  87. else:
  88. if block:
  89. deadline = monotonic() + timeout
  90. if not self._rlock.acquire(block, timeout):
  91. raise Empty
  92. try:
  93. if block:
  94. timeout = deadline - monotonic()
  95. if timeout < 0 or not self._poll(timeout):
  96. raise Empty
  97. elif not self._poll():
  98. raise Empty
  99. res = self._recv_bytes()
  100. self._sem.release()
  101. finally:
  102. self._rlock.release()
  103. # unserialize the data after having released the lock
  104. return ForkingPickler.loads(res)
  105. def qsize(self):
  106. # Raises NotImplementedError on macOS because
  107. # of broken sem_getvalue()
  108. return self._maxsize - self._sem._semlock._get_value()
  109. def empty(self):
  110. return not self._poll()
  111. def full(self):
  112. return self._sem._semlock._is_zero()
  113. def get_nowait(self):
  114. return self.get(False)
  115. def put_nowait(self, obj):
  116. return self.put(obj, False)
  117. def close(self):
  118. self._closed = True
  119. try:
  120. self._reader.close()
  121. finally:
  122. close = self._close
  123. if close:
  124. self._close = None
  125. close()
  126. def join_thread(self):
  127. debug('Queue.join_thread()')
  128. assert self._closed
  129. if self._jointhread:
  130. self._jointhread()
  131. def cancel_join_thread(self):
  132. debug('Queue.cancel_join_thread()')
  133. self._joincancelled = True
  134. try:
  135. self._jointhread.cancel()
  136. except AttributeError:
  137. pass
  138. def _start_thread(self):
  139. debug('Queue._start_thread()')
  140. # Start thread which transfers data from buffer to pipe
  141. self._buffer.clear()
  142. self._thread = threading.Thread(
  143. target=Queue._feed,
  144. args=(self._buffer, self._notempty, self._send_bytes,
  145. self._wlock, self._writer.close, self._ignore_epipe),
  146. name='QueueFeederThread'
  147. )
  148. self._thread.daemon = True
  149. debug('doing self._thread.start()')
  150. self._thread.start()
  151. debug('... done self._thread.start()')
  152. # On process exit we will wait for data to be flushed to pipe.
  153. #
  154. # However, if this process created the queue then all
  155. # processes which use the queue will be descendants of this
  156. # process. Therefore waiting for the queue to be flushed
  157. # is pointless once all the child processes have been joined.
  158. created_by_this_process = (self._opid == os.getpid())
  159. if not self._joincancelled and not created_by_this_process:
  160. self._jointhread = Finalize(
  161. self._thread, Queue._finalize_join,
  162. [weakref.ref(self._thread)],
  163. exitpriority=-5
  164. )
  165. # Send sentinel to the thread queue object when garbage collected
  166. self._close = Finalize(
  167. self, Queue._finalize_close,
  168. [self._buffer, self._notempty],
  169. exitpriority=10
  170. )
  171. @staticmethod
  172. def _finalize_join(twr):
  173. debug('joining queue thread')
  174. thread = twr()
  175. if thread is not None:
  176. thread.join()
  177. debug('... queue thread joined')
  178. else:
  179. debug('... queue thread already dead')
  180. @staticmethod
  181. def _finalize_close(buffer, notempty):
  182. debug('telling queue thread to quit')
  183. with notempty:
  184. buffer.append(_sentinel)
  185. notempty.notify()
  186. @staticmethod
  187. def _feed(buffer, notempty, send_bytes, writelock, close, ignore_epipe):
  188. debug('starting thread to feed data to pipe')
  189. nacquire = notempty.acquire
  190. nrelease = notempty.release
  191. nwait = notempty.wait
  192. bpopleft = buffer.popleft
  193. sentinel = _sentinel
  194. if sys.platform != 'win32':
  195. wacquire = writelock.acquire
  196. wrelease = writelock.release
  197. else:
  198. wacquire = None
  199. try:
  200. while 1:
  201. nacquire()
  202. try:
  203. if not buffer:
  204. nwait()
  205. finally:
  206. nrelease()
  207. try:
  208. while 1:
  209. obj = bpopleft()
  210. if obj is sentinel:
  211. debug('feeder thread got sentinel -- exiting')
  212. close()
  213. return
  214. # serialize the data before acquiring the lock
  215. obj = ForkingPickler.dumps(obj)
  216. if wacquire is None:
  217. send_bytes(obj)
  218. else:
  219. wacquire()
  220. try:
  221. send_bytes(obj)
  222. finally:
  223. wrelease()
  224. except IndexError:
  225. pass
  226. except Exception as exc:
  227. if ignore_epipe and get_errno(exc) == errno.EPIPE:
  228. return
  229. # Since this runs in a daemon thread the resources it uses
  230. # may be become unusable while the process is cleaning up.
  231. # We ignore errors which happen after the process has
  232. # started to cleanup.
  233. try:
  234. if is_exiting():
  235. info('error in queue thread: %r', exc, exc_info=True)
  236. else:
  237. if not error('error in queue thread: %r', exc,
  238. exc_info=True):
  239. import traceback
  240. traceback.print_exc()
  241. except Exception:
  242. pass
  243. _sentinel = object()
  244. class JoinableQueue(Queue):
  245. '''
  246. A queue type which also supports join() and task_done() methods
  247. Note that if you do not call task_done() for each finished task then
  248. eventually the counter's semaphore may overflow causing Bad Things
  249. to happen.
  250. '''
  251. def __init__(self, maxsize=0, *args, **kwargs):
  252. try:
  253. ctx = kwargs['ctx']
  254. except KeyError:
  255. raise TypeError('missing 1 required keyword argument: ctx')
  256. Queue.__init__(self, maxsize, ctx=ctx)
  257. self._unfinished_tasks = ctx.Semaphore(0)
  258. self._cond = ctx.Condition()
  259. def __getstate__(self):
  260. return Queue.__getstate__(self) + (self._cond, self._unfinished_tasks)
  261. def __setstate__(self, state):
  262. Queue.__setstate__(self, state[:-2])
  263. self._cond, self._unfinished_tasks = state[-2:]
  264. def put(self, obj, block=True, timeout=None):
  265. assert not self._closed
  266. if not self._sem.acquire(block, timeout):
  267. raise Full
  268. with self._notempty:
  269. with self._cond:
  270. if self._thread is None:
  271. self._start_thread()
  272. self._buffer.append(obj)
  273. self._unfinished_tasks.release()
  274. self._notempty.notify()
  275. def task_done(self):
  276. with self._cond:
  277. if not self._unfinished_tasks.acquire(False):
  278. raise ValueError('task_done() called too many times')
  279. if self._unfinished_tasks._semlock._is_zero():
  280. self._cond.notify_all()
  281. def join(self):
  282. with self._cond:
  283. if not self._unfinished_tasks._semlock._is_zero():
  284. self._cond.wait()
  285. class _SimpleQueue(object):
  286. '''
  287. Simplified Queue type -- really just a locked pipe
  288. '''
  289. def __init__(self, rnonblock=False, wnonblock=False, ctx=None):
  290. self._reader, self._writer = connection.Pipe(
  291. duplex=False, rnonblock=rnonblock, wnonblock=wnonblock,
  292. )
  293. self._poll = self._reader.poll
  294. self._rlock = self._wlock = None
  295. def empty(self):
  296. return not self._poll()
  297. def __getstate__(self):
  298. context.assert_spawning(self)
  299. return (self._reader, self._writer, self._rlock, self._wlock)
  300. def __setstate__(self, state):
  301. (self._reader, self._writer, self._rlock, self._wlock) = state
  302. def get_payload(self):
  303. return self._reader.recv_bytes()
  304. def send_payload(self, value):
  305. self._writer.send_bytes(value)
  306. def get(self):
  307. # unserialize the data after having released the lock
  308. return ForkingPickler.loads(self.get_payload())
  309. def put(self, obj):
  310. # serialize the data before acquiring the lock
  311. self.send_payload(ForkingPickler.dumps(obj))
  312. class SimpleQueue(_SimpleQueue):
  313. def __init__(self, *args, **kwargs):
  314. try:
  315. ctx = kwargs['ctx']
  316. except KeyError:
  317. raise TypeError('missing required keyword argument: ctx')
  318. self._reader, self._writer = connection.Pipe(duplex=False)
  319. self._rlock = ctx.Lock()
  320. self._wlock = ctx.Lock() if sys.platform != 'win32' else None
  321. def get_payload(self):
  322. with self._rlock:
  323. return self._reader.recv_bytes()
  324. def send_payload(self, value):
  325. if self._wlock is None:
  326. # writes to a message oriented win32 pipe are atomic
  327. self._writer.send_bytes(value)
  328. else:
  329. with self._wlock:
  330. self._writer.send_bytes(value)