Use semaphore to synchronize threads

Using `threading.Semaphore` makes it easier to cap the number of concurrently ran tasks. It also makes it possible to remove busy wait in child thread by waiting for semaphore. Also I've updated code to use the backpressure pattern - the new tasks are scheduled as soon as the user consumes the old ones.
python · obaltian · Jul 20, 2025 · Jul 20, 2025 · Jul 21, 2025 · Jul 21, 2025
commit f6f423cad66cec8901a51f47ea8570b01a525310
@@ -14,6 +14,7 @@
 #
 
 import collections
+import functools
 import itertools
 import os
 import queue
@@ -395,32 +396,20 @@ def _guarded_task_generation(self, result_job, func, iterable):
             yield (result_job, i+1, _helper_reraises_exception, (e,), {})
 
     def _guarded_task_generation_lazy(self, result_job, func, iterable,
-                                      lazy_task_gen_helper):
-        '''Provides a generator of tasks for imap and imap_unordered with
+                                      backpressure_sema):
+        """Provides a generator of tasks for imap and imap_unordered with
         appropriate handling for iterables which throw exceptions during
-        iteration.'''
-        if not lazy_task_gen_helper.feature_enabled:
-            yield from self._guarded_task_generation(result_job, func, iterable)
-            return
-
+        iteration."""
         try:
             i = -1
             enumerated_iter = iter(enumerate(iterable))
-            thread = threading.current_thread()
-            max_generated_tasks = self._processes + lazy_task_gen_helper.buffersize
-
-            while thread._state == RUN:
-                with lazy_task_gen_helper.iterator_cond:
-                    if lazy_task_gen_helper.not_finished_tasks >= max_generated_tasks:
-                        continue  # wait for some task to be (picked up and) finished
-
+            while True:
+                backpressure_sema.acquire()
                 try:
-                    i, x = enumerated_iter.__next__()
+                    i, x = next(enumerated_iter)
                 except StopIteration:
                     break
-
                 yield (result_job, i, func, (x,), {})
-                lazy_task_gen_helper.tasks_generated += 1
 
         except Exception as e:
             yield (result_job, i+1, _helper_reraises_exception, (e,), {})
@@ -430,31 +419,32 @@ def imap(self, func, iterable, chunksize=1, buffersize=None):
         Equivalent of `map()` -- can be MUCH slower than `Pool.map()`.
         '''
         self._check_running()
+        if chunksize < 1:
+            raise ValueError("Chunksize must be 1+, not {0:n}".format(chunksize))
+
+        result = IMapIterator(self, buffersize)
+
+        if result._backpressure_sema is None:
+            task_generation = self._guarded_task_generation
+        else:
+            task_generation = functools.partial(
+                self._guarded_task_generation_lazy,
+                backpressure_sema=result._backpressure_sema,
+            )
+
         if chunksize == 1:
-            result = IMapIterator(self, buffersize)
             self._taskqueue.put(
                 (
-                    self._guarded_task_generation_lazy(result._job,
-                                                       func,
-                                                       iterable,
-                                                       result._lazy_task_gen_helper),
+                    task_generation(result._job, func, iterable),
                     result._set_length,
                 )
             )
             return result
         else:
-            if chunksize < 1:
-                raise ValueError(
-                    "Chunksize must be 1+, not {0:n}".format(
-                        chunksize))
             task_batches = Pool._get_tasks(func, iterable, chunksize)
-            result = IMapIterator(self, buffersize)
             self._taskqueue.put(
                 (
-                    self._guarded_task_generation_lazy(result._job,
-                                                       mapstar,
-                                                       task_batches,
-                                                       result._lazy_task_gen_helper),
+                    task_generation(result._job, mapstar, task_batches),
                     result._set_length,
                 )
             )
@@ -465,30 +455,34 @@ def imap_unordered(self, func, iterable, chunksize=1, buffersize=None):
         Like `imap()` method but ordering of results is arbitrary.
         '''
         self._check_running()
+        if chunksize < 1:
+            raise ValueError(
+                "Chunksize must be 1+, not {0!r}".format(chunksize)
+            )
+
+        result = IMapUnorderedIterator(self, buffersize)
+
+        if result._backpressure_sema is None:
+            task_generation = self._guarded_task_generation
+        else:
+            task_generation = functools.partial(
+                self._guarded_task_generation_lazy,
+                backpressure_sema=result._backpressure_sema,
+            )
+
         if chunksize == 1:
-            result = IMapUnorderedIterator(self, buffersize)
             self._taskqueue.put(
                 (
-                    self._guarded_task_generation_lazy(result._job,
-                                                       func,
-                                                       iterable,
-                                                       result._lazy_task_gen_helper),
+                    task_generation(result._job, func, iterable),
                     result._set_length,
                 )
             )
             return result
         else:
-            if chunksize < 1:
-                raise ValueError(
-                    "Chunksize must be 1+, not {0!r}".format(chunksize))
             task_batches = Pool._get_tasks(func, iterable, chunksize)
-            result = IMapUnorderedIterator(self, buffersize)
             self._taskqueue.put(
                 (
-                    self._guarded_task_generation_lazy(result._job,
-                                                       mapstar,
-                                                       task_batches,
-                                                       result._lazy_task_gen_helper),
+                    task_generation(result._job, mapstar, task_batches),
                     result._set_length,
                 )
             )
@@ -889,7 +883,13 @@ def __init__(self, pool, buffersize):
         self._length = None
         self._unsorted = {}
         self._cache[self._job] = self
-        self._lazy_task_gen_helper = _LazyTaskGenHelper(buffersize, self._cond)
+
+        if buffersize is None:
+            self._backpressure_sema = None
+        else:
+            self._backpressure_sema = threading.Semaphore(
+                value=self._pool._processes + buffersize
+            )
 
     def __iter__(self):
         return self
@@ -910,7 +910,9 @@ def next(self, timeout=None):
                         self._pool = None
                         raise StopIteration from None
                     raise TimeoutError from None
-            self._lazy_task_gen_helper.tasks_finished += 1
+
+        if self._backpressure_sema:
+            self._backpressure_sema.release()
 
         success, value = item
         if success:
@@ -959,22 +961,6 @@ def _set(self, i, obj):
                 del self._cache[self._job]
                 self._pool = None
 
-#
-# Class to store stats for lazy task generation and share them
-# between the main thread and `_guarded_task_generation()` thread.
-#
-class _LazyTaskGenHelper(object):
-    def __init__(self, buffersize, iterator_cond):
-        self.feature_enabled = buffersize is not None
-        self.buffersize = buffersize
-        self.tasks_generated = 0
-        self.tasks_finished = 0
-        self.iterator_cond = iterator_cond
-
-    @property
-    def not_finished_tasks(self):
-        return self.tasks_generated - self.tasks_finished
-
 #
 #
 #