queue.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344
  1. #
  2. # Module implementing queues
  3. #
  4. # processing/queue.py
  5. #
  6. # Copyright (c) 2006-2008, R Oudkerk --- see COPYING.txt
  7. #
  8. __all__ = ['Queue', 'SimpleQueue']
  9. import sys
  10. import os
  11. import threading
  12. import collections
  13. import time
  14. import atexit
  15. import weakref
  16. from Queue import Empty, Full
  17. from processing import _processing, Pipe, currentProcess
  18. from processing.synchronize import Lock, BoundedSemaphore
  19. from processing.logger import debug, subWarning
  20. from processing.finalize import Finalize
  21. from processing.process import _exitFunction, _registerAfterFork
  22. from processing.forking import assertSpawning
  23. #
  24. # Cleanup function of `processing` should run before that of `threading`
  25. #
  26. atexit._exithandlers.remove((_exitFunction, (), {}))
  27. atexit._exithandlers.append((_exitFunction, (), {}))
  28. #
  29. # Queue type using a pipe, buffer and thread
  30. #
  31. class Queue(object):
  32. def __init__(self, maxsize=0):
  33. reader, writer = Pipe(duplex=False)
  34. rlock = Lock()
  35. if sys.platform == 'win32':
  36. wlock = None
  37. else:
  38. wlock = Lock()
  39. if maxsize < 0:
  40. maxsize = 0
  41. if maxsize == 0:
  42. sem = None
  43. else:
  44. sem = BoundedSemaphore(maxsize)
  45. state = maxsize, reader, writer, rlock, wlock, sem, os.getpid()
  46. self.__setstate__(state)
  47. if sys.platform != 'win32':
  48. _registerAfterFork(self, Queue._afterFork)
  49. def __getstate__(self):
  50. assertSpawning(self)
  51. return self._state
  52. def __setstate__(self, state):
  53. (self._maxsize, self._reader, self._writer,
  54. self._rlock, self._wlock, self._sem, self._opid) = self._state = state
  55. self._send = self._writer.send
  56. self._recv = self._reader.recv
  57. self._poll = self._reader.poll
  58. self._afterFork()
  59. def _afterFork(self):
  60. debug('Queue._afterFork()')
  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. def put(self, obj, block=True, timeout=None):
  69. assert not self._closed
  70. if self._sem is not None:
  71. if not self._sem.acquire(block, timeout):
  72. raise Full
  73. self._notempty.acquire()
  74. try:
  75. if self._thread is None:
  76. self._startThread()
  77. self._buffer.append(obj)
  78. self._notempty.notify()
  79. finally:
  80. self._notempty.release()
  81. def putMany(self, iterable):
  82. assert not self._closed
  83. assert self._maxsize == 0
  84. self._notempty.acquire()
  85. try:
  86. if self._thread is None:
  87. self._startThread()
  88. self._buffer.extend(iterable)
  89. self._notempty.notify()
  90. finally:
  91. self._notempty.release()
  92. def get(self, block=True, timeout=None):
  93. if block and timeout is None:
  94. self._rlock.acquire()
  95. try:
  96. res = self._recv()
  97. if self._sem:
  98. self._sem.release()
  99. return res
  100. finally:
  101. self._rlock.release()
  102. else:
  103. if block:
  104. deadline = time.time() + timeout
  105. if not self._rlock.acquire(block, timeout):
  106. raise Empty
  107. try:
  108. if not self._poll(block and (deadline-time.time()) or 0.0):
  109. raise Empty
  110. res = self._recv()
  111. if self._sem:
  112. self._sem.release()
  113. return res
  114. finally:
  115. self._rlock.release()
  116. def empty(self):
  117. # Even more unreliable than Queue.Queue.empty(): True can be
  118. # returned when enqueued items are buffered but none are
  119. # yet in the pipe
  120. return not self._poll()
  121. def full(self):
  122. if self._sem:
  123. if self._sem.acquire(False):
  124. self._sem.release()
  125. return False
  126. return True
  127. else:
  128. return False
  129. def getNoWait(self):
  130. return self.get(False)
  131. def putNoWait(self, obj):
  132. return self.put(obj, False)
  133. def close(self):
  134. self._closed = True
  135. self._reader.close()
  136. if self._close:
  137. self._close()
  138. def joinThread(self):
  139. debug('Queue.joinThread()')
  140. assert self._closed
  141. if self._jointhread:
  142. self._jointhread()
  143. def cancelJoin(self):
  144. debug('Queue.cancelJoin()')
  145. self._joincancelled = True
  146. try:
  147. self._jointhread.cancel()
  148. except AttributeError:
  149. pass
  150. def _startThread(self):
  151. debug('Queue._startThread()')
  152. # Start thread which transfers data from buffer to pipe
  153. self._buffer.clear()
  154. self._thread = threading.Thread(
  155. target=Queue._feed,
  156. args=(self._buffer, self._notempty, self._send,
  157. self._wlock, self._writer.close),
  158. name='QueueFeederThread'
  159. )
  160. self._thread.setDaemon(True)
  161. debug('doing self._thread.start()')
  162. self._thread.start()
  163. debug('... done self._thread.start()')
  164. # On process exit we will wait for data to be flushed to pipe.
  165. #
  166. # However, if this process created the queue then all
  167. # processes which use the queue will be descendants of this
  168. # process. Therefore waiting for the queue to be flushed
  169. # is pointless once all the child processes have been joined.
  170. created_by_this_process = (self._opid == os.getpid())
  171. if not self._joincancelled and not created_by_this_process:
  172. self._jointhread = Finalize(
  173. self._thread, Queue._finalizeJoin,
  174. [weakref.ref(self._thread)],
  175. exitpriority=-5
  176. )
  177. # Send sentinel to the thread queue object when garbage collected
  178. self._close = Finalize(
  179. self, Queue._finalizeClose,
  180. [self._buffer, self._notempty],
  181. exitpriority=10
  182. )
  183. @staticmethod
  184. def _finalizeJoin(twr):
  185. debug('joining queue thread')
  186. thread = twr()
  187. if thread is not None:
  188. thread.join()
  189. debug('... queue thread joined')
  190. else:
  191. debug('... queue thread already dead')
  192. @staticmethod
  193. def _finalizeClose(buffer, notempty):
  194. debug('telling queue thread to quit')
  195. notempty.acquire()
  196. try:
  197. buffer.append(_sentinel)
  198. notempty.notify()
  199. finally:
  200. notempty.release()
  201. @staticmethod
  202. def _feed(buffer, notempty, send, writelock, close):
  203. debug('starting thread to feed data to pipe')
  204. nacquire = notempty.acquire
  205. nrelease = notempty.release
  206. nwait = notempty.wait
  207. bpopleft = buffer.popleft
  208. sentinel = _sentinel
  209. if sys.platform != 'win32':
  210. wacquire = writelock.acquire
  211. wrelease = writelock.release
  212. else:
  213. wacquire = None
  214. try:
  215. while 1:
  216. nacquire()
  217. try:
  218. if not buffer:
  219. nwait()
  220. finally:
  221. nrelease()
  222. try:
  223. while 1:
  224. obj = bpopleft()
  225. if obj is sentinel:
  226. debug('feeder thread got sentinel -- exiting')
  227. close()
  228. return
  229. if wacquire is None:
  230. send(obj)
  231. else:
  232. wacquire()
  233. try:
  234. send(obj)
  235. finally:
  236. wrelease()
  237. except IndexError:
  238. pass
  239. except Exception, e:
  240. # Since this runs in a daemon thread the objects it uses
  241. # may be become unusable while the process is cleaning up.
  242. # We ignore errors which happen after the process has
  243. # started to cleanup.
  244. if currentProcess()._exiting:
  245. subWarning('error in queue thread: %s', e)
  246. else:
  247. raise
  248. get_nowait = getNoWait
  249. put_nowait = putNoWait
  250. # deprecated
  251. putmany = putMany
  252. jointhread = joinThread
  253. canceljoin = cancelJoin
  254. _sentinel = object()
  255. #
  256. # Simplified Queue type -- really just a locked pipe
  257. #
  258. class SimpleQueue(object):
  259. def __init__(self):
  260. reader, writer = Pipe(duplex=False)
  261. if sys.platform == 'win32':
  262. state = reader, writer, Lock(), None
  263. else:
  264. state = reader, writer, Lock(), Lock()
  265. self.__setstate__(state)
  266. def empty(self):
  267. return not self._reader.poll()
  268. def __getstate__(self):
  269. assertSpawning(self)
  270. return self._state
  271. def __setstate__(self, state):
  272. (self._reader, self._writer, self._rlock, self._wlock) \
  273. = self._state = state
  274. recv = self._reader.recv
  275. racquire, rrelease = self._rlock.acquire, self._rlock.release
  276. def get():
  277. racquire()
  278. try:
  279. return recv()
  280. finally:
  281. rrelease()
  282. self.get = get
  283. if self._wlock is None:
  284. # writes to a message oriented win32 pipe are atomic
  285. self.put = self._writer.send
  286. else:
  287. send = self._writer.send
  288. wacquire, wrelease = self._wlock.acquire, self._wlock.release
  289. def put(obj):
  290. wacquire()
  291. try:
  292. return send(obj)
  293. finally:
  294. wrelease()
  295. self.put = put