Skip to content

Commit

Permalink
refactor: set updateDelays
Browse files Browse the repository at this point in the history
  • Loading branch information
roggervalf committed Dec 18, 2023
1 parent fbf757d commit 20b82e5
Showing 1 changed file with 45 additions and 25 deletions.
70 changes: 45 additions & 25 deletions python/bullmq/worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,30 +64,46 @@ async def run(self):
token_postfix = 0

while not self.closed:
while not self.waiting and len(self.processing) < self.opts.get("concurrency") and not self.closing and not self.limitUntil:
num_total = len(self.processing)
while not self.waiting and num_total < self.opts.get("concurrency") and not self.closing and (not self.limitUntil or num_total == 0):
token_postfix+=1
token = f'{self.id}:{token_postfix}'
waiting_job = asyncio.ensure_future(self.getNextJob(token))
self.processing.add(waiting_job)

try:
jobs, pending = await getCompleted(self.processing)

jobs_to_process = [self.processJob(job, job.token) for job in jobs]
processing_jobs = [asyncio.ensure_future(
j) for j in jobs_to_process]
pending.update(processing_jobs)
self.processing = pending
num_total = len(self.processing)

if (len(jobs) == 0 or len(self.processing) == 0) and self.closing:
# We are done processing so we can close the queue
if self.waiting and num_total > 1:
# we have a job waiting but we have others that we could start processing already
break

except Exception as e:
# This should never happen or we will have an endless loop
print("ERROR:", e)
traceback.print_exc()
return
try:
jobs, pending = await getCompleted(self.processing)

jobs_to_process = [self.processJob(job, job.token) for job in jobs]
processing_jobs = [asyncio.ensure_future(
j) for j in jobs_to_process]
pending.update(processing_jobs)
self.processing = pending

# no more jobs waiting but we have others that could start processing already
if (len(jobs)==0 and num_total>1):
break

# if there are potential jobs to be processed and blockUntil is set, we should exit to avoid waiting
# for processing this job.
if self.blockUntil:
break

except Exception as e:
# This should never happen or we will have an endless loop
print("ERROR:", e)
traceback.print_exc()
return

if (len(jobs) == 0 or len(self.processing) == 0) and self.closing:
# We are done processing so we can close the queue
break

self.running = False
self.timer.stop()
Expand Down Expand Up @@ -119,6 +135,11 @@ async def moveToActive(self, token: str, job_id: str = None):
if job_id and job_id.startswith('0:'):
self.blockUntil = int(job_id.split(':')[1]) or 0

# remove marker from active list
await self.client.lrem(self.scripts.toKey('active'), 1, job_id)
if self.blockUntil > 0:
return

result = await self.scripts.moveToActive(token, self.opts, job_id)
job_data = None
id = None
Expand All @@ -127,21 +148,15 @@ async def moveToActive(self, token: str, job_id: str = None):

if result:
job_data, id, limit_until, delay_until = result
self.updateDelays(limit_until, delay_until)
return self.nextJobFromJobData(job_data, id, token)

return self.nextJobFromJobData(job_data, id, limit_until, delay_until, token)

def nextJobFromJobData(self, job_data = None, job_id: str = None, limit_until: int = 0,
delay_until: int = 0, token: str = None):
self.limitUntil = max(limit_until, 0) or 0

def nextJobFromJobData(self, job_data = None, job_id: str = None, token: str = None):
if not job_data:
if not self.drained:
self.drained = True
self.blockUntil = 0

if delay_until:
self.blockUntil = max(delay_until, 0) or 0

if job_data:
self.drained = False
job_instance = Job.fromJSON(self, job_data, job_id)
Expand All @@ -159,12 +174,17 @@ async def waitForJob(self):
job_id = await self.bclient.brpoplpush(self.scripts.keys["wait"], self.scripts.keys["active"], timeout)

return job_id

def updateDelays(self, limit_until = 0, delay_until = 0):
self.limitUntil = max(limit_until, 0) or 0
self.blockUntil = max(delay_until, 0) or 0

async def processJob(self, job: Job, token: str):
try:
self.jobs.add((job, token))
result = await self.processor(job, token)
if not self.forceClosing:
# TODO get next job from this script
await self.scripts.moveToCompleted(job, result, job.opts.get("removeOnComplete", False), token, self.opts, fetchNext=not self.closing)
job.returnvalue = result
self.emit("completed", job, result)
Expand Down

0 comments on commit 20b82e5

Please sign in to comment.