Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 10 additions & 7 deletions Lib/multiprocessing/pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -711,8 +711,9 @@ def _terminate_pool(cls, taskqueue, inqueue, outqueue, pool, change_notifier,
change_notifier.put(None)

task_handler._state = TERMINATE
# Release all semaphores to wake up task_handler to stop.
for buffersize_sema in tuple(taskqueue_buffersize_semaphores):
# Snapshot the set before releasing semaphores: the result handler can
# remove completed ones concurrently.
for buffersize_sema in list(taskqueue_buffersize_semaphores):
buffersize_sema.release()
taskqueue_buffersize_semaphores.discard(buffersize_sema)

Expand Down Expand Up @@ -935,11 +936,13 @@ def next(self, timeout=None):
return value
raise value

def _stop_iterator(self):
def _release_pool(self):
if self._pool is not None:
# `self._pool` could be set to `None` in previous `.next()` calls
self._pool._taskqueue_buffersize_semaphores.discard(self._buffersize_sema)
self._pool = None

def _stop_iterator(self):
self._release_pool()
raise StopIteration from None

__next__ = next # XXX
Expand All @@ -959,15 +962,15 @@ def _set(self, i, obj):

if self._index == self._length:
del self._cache[self._job]
self._pool = None
self._release_pool()

def _set_length(self, length):
with self._cond:
self._length = length
if self._index == self._length:
self._cond.notify()
del self._cache[self._job]
self._pool = None
self._release_pool()

#
# Class whose instances are returned by `Pool.imap_unordered()`
Expand All @@ -982,7 +985,7 @@ def _set(self, i, obj):
self._cond.notify()
if self._index == self._length:
del self._cache[self._job]
self._pool = None
self._release_pool()

#
#
Expand Down
21 changes: 21 additions & 0 deletions Lib/test/_test_multiprocessing.py
Original file line number Diff line number Diff line change
Expand Up @@ -3276,6 +3276,27 @@ def test_imap_and_imap_unordered_with_buffersize_on_empty_iterable(
self.assertIsNone(next(res, None))
self.assertIsNone(next(res, None))

@warnings_helper.ignore_fork_in_thread_deprecation_warnings()
def test_imap_with_buffersize_releases_semaphores(self):
if self.TYPE == 'manager':
self.skipTest('test requires access to the pool')

p = self.Pool(1)
try:
for method_name in ('imap', 'imap_unordered'):
for chunksize in (1, 2):
for values in ([], [1, 2, 3]):
with self.subTest(method=method_name,
chunksize=chunksize, values=values):
method = getattr(p, method_name)
result = method(str, values, chunksize=chunksize,
buffersize=1)
self.assertCountEqual(result, map(str, values))
self.assertFalse(p._taskqueue_buffersize_semaphores)
finally:
p.terminate()
p.join()

@warnings_helper.ignore_fork_in_thread_deprecation_warnings()
def test_imap_with_buffersize_on_infinite_iterable(self):
if self.TYPE != "threads":
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
Fix a memory leak in :meth:`multiprocessing.pool.Pool.imap` and
:meth:`~multiprocessing.pool.Pool.imap_unordered` when the *buffersize*
parameter is used.
Loading