Fix result_generator hanging when the job generator produces nothing.

This commit is contained in:
Ethan Dalool 2020-11-09 11:03:08 -08:00
parent cbf546ee69
commit 63689e02c0

View file

@ -255,12 +255,18 @@ class ThreadPool:
# the threads can keep waking up and seeing no more jobs. # the threads can keep waking up and seeing no more jobs.
if not self.closed: if not self.closed:
self._jobs_available.clear() self._jobs_available.clear()
if self._result_queue is not None:
# If the user provided a generator to add_generator that
# actually produces no items, and then immediately starts
# waiting inside result_generator for the results, they
# will hang as _result_queue never gets anything.
# So, here's this.
self._result_queue.put(NO_MORE_JOBS)
return NO_MORE_JOBS return NO_MORE_JOBS
else:
if self._result_queue is not None: if self._result_queue is not None:
# This will block if the queue is full. # This will block if the queue is full.
self._result_queue.put(job) self._result_queue.put(job)
return job return job
def join(self): def join(self):
@ -328,6 +334,9 @@ class ThreadPool:
# jobs. # jobs.
while self._jobs_available.is_set() or not self._result_queue.empty(): while self._jobs_available.is_set() or not self._result_queue.empty():
job = self._result_queue.get() job = self._result_queue.get()
if job is NO_MORE_JOBS:
self._result_queue.task_done()
break
job.join() job.join()
yield job yield job
self._result_queue.task_done() self._result_queue.task_done()