blob: f6f02b66656bd453252b644389e6ffc0f6dbea9c [file] [log] [blame]
Benjamin Petersone711caf2008-06-11 16:44:04 +00001#
2# Module implementing queues
3#
4# multiprocessing/queues.py
5#
R. David Murray3fc969a2010-12-14 01:38:16 +00006# Copyright (c) 2006-2008, R Oudkerk
Richard Oudkerk3e268aa2012-04-30 12:13:55 +01007# Licensed to PSF under a Contributor Agreement.
Benjamin Petersone711caf2008-06-11 16:44:04 +00008#
9
Jesse Noller14f3ae22009-03-31 03:37:07 +000010__all__ = ['Queue', 'SimpleQueue', 'JoinableQueue']
Benjamin Petersone711caf2008-06-11 16:44:04 +000011
12import sys
13import os
14import threading
15import collections
16import time
Benjamin Petersone711caf2008-06-11 16:44:04 +000017import weakref
Antoine Pitroudc19c242011-07-16 01:51:58 +020018import errno
Benjamin Petersone711caf2008-06-11 16:44:04 +000019
20from queue import Empty, Full
21import _multiprocessing
Antoine Pitroubdb1cf12012-03-05 19:28:37 +010022from multiprocessing.connection import Pipe
Benjamin Petersone711caf2008-06-11 16:44:04 +000023from multiprocessing.synchronize import Lock, BoundedSemaphore, Semaphore, Condition
24from multiprocessing.util import debug, info, Finalize, register_after_fork
25from multiprocessing.forking import assert_spawning
26
27#
28# Queue type using a pipe, buffer and thread
29#
30
31class Queue(object):
32
33 def __init__(self, maxsize=0):
34 if maxsize <= 0:
35 maxsize = _multiprocessing.SemLock.SEM_VALUE_MAX
36 self._maxsize = maxsize
37 self._reader, self._writer = Pipe(duplex=False)
38 self._rlock = Lock()
39 self._opid = os.getpid()
40 if sys.platform == 'win32':
41 self._wlock = None
42 else:
43 self._wlock = Lock()
44 self._sem = BoundedSemaphore(maxsize)
Antoine Pitroudc19c242011-07-16 01:51:58 +020045 # For use by concurrent.futures
46 self._ignore_epipe = False
Benjamin Petersone711caf2008-06-11 16:44:04 +000047
48 self._after_fork()
49
50 if sys.platform != 'win32':
51 register_after_fork(self, Queue._after_fork)
52
53 def __getstate__(self):
54 assert_spawning(self)
Antoine Pitroufb960892011-07-20 02:01:39 +020055 return (self._ignore_epipe, self._maxsize, self._reader, self._writer,
Benjamin Petersone711caf2008-06-11 16:44:04 +000056 self._rlock, self._wlock, self._sem, self._opid)
57
58 def __setstate__(self, state):
Antoine Pitroufb960892011-07-20 02:01:39 +020059 (self._ignore_epipe, self._maxsize, self._reader, self._writer,
Benjamin Petersone711caf2008-06-11 16:44:04 +000060 self._rlock, self._wlock, self._sem, self._opid) = state
61 self._after_fork()
62
63 def _after_fork(self):
64 debug('Queue._after_fork()')
65 self._notempty = threading.Condition(threading.Lock())
66 self._buffer = collections.deque()
67 self._thread = None
68 self._jointhread = None
69 self._joincancelled = False
70 self._closed = False
71 self._close = None
72 self._send = self._writer.send
73 self._recv = self._reader.recv
74 self._poll = self._reader.poll
75
76 def put(self, obj, block=True, timeout=None):
77 assert not self._closed
78 if not self._sem.acquire(block, timeout):
79 raise Full
80
81 self._notempty.acquire()
82 try:
83 if self._thread is None:
84 self._start_thread()
85 self._buffer.append(obj)
86 self._notempty.notify()
87 finally:
88 self._notempty.release()
89
90 def get(self, block=True, timeout=None):
91 if block and timeout is None:
92 self._rlock.acquire()
93 try:
94 res = self._recv()
95 self._sem.release()
96 return res
97 finally:
98 self._rlock.release()
99
100 else:
101 if block:
102 deadline = time.time() + timeout
103 if not self._rlock.acquire(block, timeout):
104 raise Empty
105 try:
Antoine Pitroua3651132011-11-10 00:37:09 +0100106 if block:
107 timeout = deadline - time.time()
108 if timeout < 0 or not self._poll(timeout):
109 raise Empty
110 elif not self._poll():
Benjamin Petersone711caf2008-06-11 16:44:04 +0000111 raise Empty
112 res = self._recv()
113 self._sem.release()
114 return res
115 finally:
116 self._rlock.release()
117
118 def qsize(self):
Benjamin Peterson87c8d872009-06-11 22:54:11 +0000119 # Raises NotImplementedError on Mac OSX because of broken sem_getvalue()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000120 return self._maxsize - self._sem._semlock._get_value()
121
122 def empty(self):
123 return not self._poll()
124
125 def full(self):
126 return self._sem._semlock._is_zero()
127
128 def get_nowait(self):
129 return self.get(False)
130
131 def put_nowait(self, obj):
132 return self.put(obj, False)
133
134 def close(self):
135 self._closed = True
136 self._reader.close()
137 if self._close:
138 self._close()
139
140 def join_thread(self):
141 debug('Queue.join_thread()')
142 assert self._closed
143 if self._jointhread:
144 self._jointhread()
145
146 def cancel_join_thread(self):
147 debug('Queue.cancel_join_thread()')
148 self._joincancelled = True
149 try:
150 self._jointhread.cancel()
151 except AttributeError:
152 pass
153
154 def _start_thread(self):
155 debug('Queue._start_thread()')
156
157 # Start thread which transfers data from buffer to pipe
158 self._buffer.clear()
159 self._thread = threading.Thread(
160 target=Queue._feed,
161 args=(self._buffer, self._notempty, self._send,
Antoine Pitroudc19c242011-07-16 01:51:58 +0200162 self._wlock, self._writer.close, self._ignore_epipe),
Benjamin Petersone711caf2008-06-11 16:44:04 +0000163 name='QueueFeederThread'
164 )
Benjamin Peterson72753702008-08-18 18:09:21 +0000165 self._thread.daemon = True
Benjamin Petersone711caf2008-06-11 16:44:04 +0000166
167 debug('doing self._thread.start()')
168 self._thread.start()
169 debug('... done self._thread.start()')
170
171 # On process exit we will wait for data to be flushed to pipe.
172 #
173 # However, if this process created the queue then all
174 # processes which use the queue will be descendants of this
175 # process. Therefore waiting for the queue to be flushed
176 # is pointless once all the child processes have been joined.
177 created_by_this_process = (self._opid == os.getpid())
178 if not self._joincancelled and not created_by_this_process:
179 self._jointhread = Finalize(
180 self._thread, Queue._finalize_join,
181 [weakref.ref(self._thread)],
182 exitpriority=-5
183 )
184
185 # Send sentinel to the thread queue object when garbage collected
186 self._close = Finalize(
187 self, Queue._finalize_close,
188 [self._buffer, self._notempty],
189 exitpriority=10
190 )
191
192 @staticmethod
193 def _finalize_join(twr):
194 debug('joining queue thread')
195 thread = twr()
196 if thread is not None:
197 thread.join()
198 debug('... queue thread joined')
199 else:
200 debug('... queue thread already dead')
201
202 @staticmethod
203 def _finalize_close(buffer, notempty):
204 debug('telling queue thread to quit')
205 notempty.acquire()
206 try:
207 buffer.append(_sentinel)
208 notempty.notify()
209 finally:
210 notempty.release()
211
212 @staticmethod
Antoine Pitroudc19c242011-07-16 01:51:58 +0200213 def _feed(buffer, notempty, send, writelock, close, ignore_epipe):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000214 debug('starting thread to feed data to pipe')
215 from .util import is_exiting
216
217 nacquire = notempty.acquire
218 nrelease = notempty.release
219 nwait = notempty.wait
220 bpopleft = buffer.popleft
221 sentinel = _sentinel
222 if sys.platform != 'win32':
223 wacquire = writelock.acquire
224 wrelease = writelock.release
225 else:
226 wacquire = None
227
228 try:
229 while 1:
230 nacquire()
231 try:
232 if not buffer:
233 nwait()
234 finally:
235 nrelease()
236 try:
237 while 1:
238 obj = bpopleft()
239 if obj is sentinel:
240 debug('feeder thread got sentinel -- exiting')
241 close()
242 return
243
244 if wacquire is None:
245 send(obj)
Andrew Svetlov6b973742012-11-03 15:36:01 +0200246 # Delete references to object. See issue16284
247 del obj
Benjamin Petersone711caf2008-06-11 16:44:04 +0000248 else:
249 wacquire()
250 try:
251 send(obj)
Andrew Svetlov6b973742012-11-03 15:36:01 +0200252 # Delete references to object. See issue16284
253 del obj
Benjamin Petersone711caf2008-06-11 16:44:04 +0000254 finally:
255 wrelease()
256 except IndexError:
257 pass
258 except Exception as e:
Antoine Pitroudc19c242011-07-16 01:51:58 +0200259 if ignore_epipe and getattr(e, 'errno', 0) == errno.EPIPE:
260 return
Benjamin Petersone711caf2008-06-11 16:44:04 +0000261 # Since this runs in a daemon thread the resources it uses
262 # may be become unusable while the process is cleaning up.
263 # We ignore errors which happen after the process has
264 # started to cleanup.
265 try:
266 if is_exiting():
267 info('error in queue thread: %s', e)
268 else:
269 import traceback
270 traceback.print_exc()
271 except Exception:
272 pass
273
274_sentinel = object()
275
276#
277# A queue type which also supports join() and task_done() methods
278#
279# Note that if you do not call task_done() for each finished task then
280# eventually the counter's semaphore may overflow causing Bad Things
281# to happen.
282#
283
284class JoinableQueue(Queue):
285
286 def __init__(self, maxsize=0):
287 Queue.__init__(self, maxsize)
288 self._unfinished_tasks = Semaphore(0)
289 self._cond = Condition()
290
291 def __getstate__(self):
292 return Queue.__getstate__(self) + (self._cond, self._unfinished_tasks)
293
294 def __setstate__(self, state):
295 Queue.__setstate__(self, state[:-2])
296 self._cond, self._unfinished_tasks = state[-2:]
297
Benjamin Peterson8719ad52009-09-11 22:24:02 +0000298 def put(self, obj, block=True, timeout=None):
299 assert not self._closed
300 if not self._sem.acquire(block, timeout):
301 raise Full
302
303 self._notempty.acquire()
304 self._cond.acquire()
305 try:
306 if self._thread is None:
307 self._start_thread()
308 self._buffer.append(obj)
309 self._unfinished_tasks.release()
310 self._notempty.notify()
311 finally:
312 self._cond.release()
313 self._notempty.release()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000314
315 def task_done(self):
316 self._cond.acquire()
317 try:
318 if not self._unfinished_tasks.acquire(False):
319 raise ValueError('task_done() called too many times')
320 if self._unfinished_tasks._semlock._is_zero():
321 self._cond.notify_all()
322 finally:
323 self._cond.release()
324
325 def join(self):
326 self._cond.acquire()
327 try:
328 if not self._unfinished_tasks._semlock._is_zero():
329 self._cond.wait()
330 finally:
331 self._cond.release()
332
333#
334# Simplified Queue type -- really just a locked pipe
335#
336
337class SimpleQueue(object):
338
339 def __init__(self):
340 self._reader, self._writer = Pipe(duplex=False)
341 self._rlock = Lock()
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100342 self._poll = self._reader.poll
Benjamin Petersone711caf2008-06-11 16:44:04 +0000343 if sys.platform == 'win32':
344 self._wlock = None
345 else:
346 self._wlock = Lock()
347 self._make_methods()
348
349 def empty(self):
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100350 return not self._poll()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000351
352 def __getstate__(self):
353 assert_spawning(self)
354 return (self._reader, self._writer, self._rlock, self._wlock)
355
356 def __setstate__(self, state):
357 (self._reader, self._writer, self._rlock, self._wlock) = state
358 self._make_methods()
359
360 def _make_methods(self):
361 recv = self._reader.recv
362 racquire, rrelease = self._rlock.acquire, self._rlock.release
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100363 def get():
Benjamin Petersone711caf2008-06-11 16:44:04 +0000364 racquire()
365 try:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100366 return recv()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000367 finally:
368 rrelease()
369 self.get = get
370
371 if self._wlock is None:
372 # writes to a message oriented win32 pipe are atomic
373 self.put = self._writer.send
374 else:
375 send = self._writer.send
376 wacquire, wrelease = self._wlock.acquire, self._wlock.release
377 def put(obj):
378 wacquire()
379 try:
380 return send(obj)
381 finally:
382 wrelease()
383 self.put = put