diff --git a/mrq/context.py b/mrq/context.py index c20d99ec..f03db67b 100644 --- a/mrq/context.py +++ b/mrq/context.py @@ -211,6 +211,10 @@ def subpool_map(pool_size, func, iterable): current_job = get_current_job() def inner_func(*args): + """ As each call to 'func' will be done in a random greenlet of the subpool, we need to + register their IDs with set_current_job() to make get_current_job() calls work properly + inside 'func'. + """ next(counter) if current_job: set_current_job(current_job) @@ -219,9 +223,20 @@ def inner_func(*args): set_current_job(None) return ret + def inner_iterable(): + """ This will be called inside the pool's main greenlet, which ID also needs to be registered """ + if current_job: + set_current_job(current_job) + + for x in iterable: + yield x + + if current_job: + set_current_job(None) + start_time = time.time() pool = gevent.pool.Pool(size=pool_size) - ret = pool.map(inner_func, iterable) + ret = pool.map(inner_func, inner_iterable()) pool.join(raise_error=True) total_time = time.time() - start_time