From e1ba6bc88d6b05d44e07ed0e87028729c6287aa9 Mon Sep 17 00:00:00 2001 From: Tom Lawton Date: Sat, 30 May 2026 13:18:31 +0100 Subject: [PATCH 1/4] Feature: add buffer_concurrency to worker for dynamic concurrency changes. Add run_worker_async_background() to return worker for control --- procrastinate/app.py | 36 ++++++++++++++++++ procrastinate/worker.py | 82 +++++++++++++++++++++++++++++++++++++++-- 2 files changed, 114 insertions(+), 4 deletions(-) diff --git a/procrastinate/app.py b/procrastinate/app.py index 982492f26..ddaf06571 100644 --- a/procrastinate/app.py +++ b/procrastinate/app.py @@ -26,6 +26,7 @@ class WorkerOptions(TypedDict): queues: NotRequired[Iterable[str]] name: NotRequired[str] concurrency: NotRequired[int] + buffer_concurrency: NotRequired[int] wait: NotRequired[bool] fetch_job_polling_interval: NotRequired[float] abort_job_polling_interval: NotRequired[float] @@ -342,6 +343,41 @@ async def f(): asyncio.run(f()) + async def run_worker_async_background(self, **kwargs: Unpack[WorkerOptions]) -> worker.Worker: + """ + Run a worker as a background asyncio task. + + Returns the worker instance immediately, allowing concurrency to be + adjusted at runtime via ``worker.set_concurrency()``. + + Parameters + ---------- + buffer_concurrency : + Number of buffered semaphore slots for future scaling. + Total semaphore capacity will be concurrency + buffer_concurrency. + install_signal_handlers : + Defaults to ``False`` for this method since the worker + runs as a task inside a larger application. + + Returns + ------- + Worker + The worker instance. Call ``worker.stop()`` for graceful shutdown. + The worker runs as ``worker.run_task`` which you can await or cancel. + """ + self.perform_import_paths() + + # Default to no signal handlers for embedded worker + if "install_signal_handlers" not in kwargs: + kwargs["install_signal_handlers"] = False + + worker = self._worker(**kwargs) + worker.run_task = asyncio.create_task( + worker.run(), + name=worker.worker_name, + ) + return worker + async def check_connection_async(self) -> bool: return await self.job_manager.check_connection_async() diff --git a/procrastinate/worker.py b/procrastinate/worker.py index 84d4a7cc3..8d9c9454b 100644 --- a/procrastinate/worker.py +++ b/procrastinate/worker.py @@ -37,6 +37,7 @@ def __init__( queues: Iterable[str] | None = None, name: str | None = WORKER_NAME, concurrency: int = WORKER_CONCURRENCY, + buffer_concurrency: int = 0, wait: bool = True, fetch_job_polling_interval: float = FETCH_JOB_POLLING_INTERVAL, abort_job_polling_interval: float = ABORT_JOB_POLLING_INTERVAL, @@ -52,6 +53,9 @@ def __init__( self.queues = queues self.worker_name = name self.concurrency = concurrency + self.buffer_concurrency = buffer_concurrency + self.total_capacity = concurrency + buffer_concurrency + self._held_buffer_slots = 0 self.wait = wait self.fetch_job_polling_interval = fetch_job_polling_interval self.abort_job_polling_interval = abort_job_polling_interval @@ -94,6 +98,57 @@ def stop(self): self._stop_event.set() + def set_concurrency(self, new_concurrency: int) -> None: + """ + Adjust the active concurrency at runtime. + + Parameters + ---------- + new_concurrency : + New concurrency level. Must be between 0 and total_capacity. + + Notes + ----- + This method is not thread-safe. Call from a single control coroutine + to avoid race conditions with the worker's internal state. + """ + if new_concurrency == self.concurrency: + return + + if new_concurrency < 0 or new_concurrency > self.total_capacity: + raise ValueError( + f"Concurrency must be between 0 and {self.total_capacity}" + ) + + old_concurrency = self.concurrency + old_buffer = self.buffer_concurrency + self.concurrency = new_concurrency + new_buffer = self.total_capacity - new_concurrency + + if new_concurrency > old_concurrency: + # Scale up: release buffer slots one at a time, decrementing counter. + # Guard against releasing more slots than actually held. + # If set_concurrency is called before _run_loop, _held_buffer_slots is 0, + # so slots_to_release becomes 0 and the loop below never runs. + # This also avoids releasing on the parent-init semaphore (capacity=concurrency) + # before _run_loop replaces it with the correct one (capacity=total_capacity). + slots_to_release = min(old_buffer - new_buffer, self._held_buffer_slots) + for _ in range(slots_to_release): + self._job_semaphore.release() + self._held_buffer_slots -= 1 + elif new_concurrency < old_concurrency: + # Scale down: no immediate action needed. + # The fetch loop will naturally hold more slots as buffer + # on the next acquisitions (since _held_buffer_slots < new_buffer). + pass + + self.logger.info( + f"Concurrency adjusted: {old_concurrency} -> {new_concurrency} " + f"(buffer: {old_buffer} -> {new_buffer})" + ) + + self.buffer_concurrency = new_buffer + async def _periodic_deferrer(self): deferrer = periodic.PeriodicDeferrer( registry=self.app.periodic_registry, @@ -327,19 +382,27 @@ async def _fetch_and_process_jobs(self): while not self._stop_event.is_set(): acquire_sem_task = asyncio.create_task(self._job_semaphore.acquire()) job = None + should_release = True try: await utils.wait_any(acquire_sem_task, self._stop_event.wait()) if self._stop_event.is_set(): break + # Check if this slot should be held as buffer + if self._held_buffer_slots < self.buffer_concurrency: + self._held_buffer_slots += 1 + should_release = False + continue + assert self.worker_id is not None job = await self.app.job_manager.fetch_job( queues=self.queues, worker_id=self.worker_id ) finally: - if (not job or self._stop_event.is_set()) and acquire_sem_task.done(): - self._job_semaphore.release() - self._new_job_event.clear() + if should_release: + if (not job or self._stop_event.is_set()) and acquire_sem_task.done(): + self._job_semaphore.release() + self._new_job_event.clear() if not job: break @@ -592,7 +655,7 @@ async def _run_loop(self): self._new_job_event.clear() self._stop_event.clear() self._running_jobs = {} - self._job_semaphore = asyncio.Semaphore(self.concurrency) + self._job_semaphore = asyncio.Semaphore(self.total_capacity) side_tasks = self._start_side_tasks() side_tasks_monitor = asyncio.create_task( self._monitor_side_tasks(side_tasks), name="side_tasks_monitor" @@ -606,6 +669,12 @@ async def _run_loop(self): try: with context: + # Hold initial buffer slots before first fetch + self._held_buffer_slots = 0 + for _ in range(self.buffer_concurrency): + await self._job_semaphore.acquire() + self._held_buffer_slots += 1 + await self._fetch_and_process_jobs() if not self.wait: self.logger.info( @@ -628,6 +697,11 @@ async def _run_loop(self): ) await self._fetch_and_process_jobs() finally: + # Release buffer slots on shutdown + for _ in range(self._held_buffer_slots): + self._job_semaphore.release() + self._held_buffer_slots -= 1 + if not side_tasks_monitor.done(): side_tasks_monitor.cancel() await self._shutdown(side_tasks=side_tasks) From dedb2368b2fd415bec2488a136fef7a8b095feb5 Mon Sep 17 00:00:00 2001 From: Tom Lawton Date: Sun, 31 May 2026 21:03:43 +0100 Subject: [PATCH 2/4] Add worker properties to check jobs and spare concurrency --- procrastinate/worker.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/procrastinate/worker.py b/procrastinate/worker.py index 8d9c9454b..96505af0a 100644 --- a/procrastinate/worker.py +++ b/procrastinate/worker.py @@ -149,6 +149,21 @@ def set_concurrency(self, new_concurrency: int) -> None: self.buffer_concurrency = new_buffer + @property + def running_jobs_count(self) -> int: + """Number of jobs currently being processed.""" + return len(self._running_jobs) + + @property + def spare_concurrency(self) -> int: + """Number of active concurrency slots currently free. + + Returns 0 when the worker is saturated (all slots occupied). + A positive value means the worker could process more jobs right now + if they were available — i.e. concurrency is NOT the bottleneck. + """ + return max(0, self.concurrency - len(self._running_jobs)) + async def _periodic_deferrer(self): deferrer = periodic.PeriodicDeferrer( registry=self.app.periodic_registry, From 01f92964ef5e283f4a6568f65f32833529a9d457 Mon Sep 17 00:00:00 2001 From: Tom Lawton Date: Wed, 3 Jun 2026 18:51:47 +0100 Subject: [PATCH 3/4] Validate concurrency inputs in Worker.__init__ --- procrastinate/worker.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/procrastinate/worker.py b/procrastinate/worker.py index 96505af0a..14d90b213 100644 --- a/procrastinate/worker.py +++ b/procrastinate/worker.py @@ -52,6 +52,14 @@ def __init__( self.app = app self.queues = queues self.worker_name = name + if concurrency < 0 or buffer_concurrency < 0: + raise ValueError( + "concurrency and buffer_concurrency must be non-negative" + ) + if concurrency + buffer_concurrency <= 0: + raise ValueError( + "concurrency + buffer_concurrency must be greater than 0" + ) self.concurrency = concurrency self.buffer_concurrency = buffer_concurrency self.total_capacity = concurrency + buffer_concurrency From 46ad5546f1281a9452c3ddcd712aeaaef61f50ae Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 3 Jun 2026 18:01:19 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- procrastinate/app.py | 4 +++- procrastinate/worker.py | 16 ++++++---------- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/procrastinate/app.py b/procrastinate/app.py index ddaf06571..f6c6a440f 100644 --- a/procrastinate/app.py +++ b/procrastinate/app.py @@ -343,7 +343,9 @@ async def f(): asyncio.run(f()) - async def run_worker_async_background(self, **kwargs: Unpack[WorkerOptions]) -> worker.Worker: + async def run_worker_async_background( + self, **kwargs: Unpack[WorkerOptions] + ) -> worker.Worker: """ Run a worker as a background asyncio task. diff --git a/procrastinate/worker.py b/procrastinate/worker.py index 14d90b213..4804e3926 100644 --- a/procrastinate/worker.py +++ b/procrastinate/worker.py @@ -53,13 +53,9 @@ def __init__( self.queues = queues self.worker_name = name if concurrency < 0 or buffer_concurrency < 0: - raise ValueError( - "concurrency and buffer_concurrency must be non-negative" - ) + raise ValueError("concurrency and buffer_concurrency must be non-negative") if concurrency + buffer_concurrency <= 0: - raise ValueError( - "concurrency + buffer_concurrency must be greater than 0" - ) + raise ValueError("concurrency + buffer_concurrency must be greater than 0") self.concurrency = concurrency self.buffer_concurrency = buffer_concurrency self.total_capacity = concurrency + buffer_concurrency @@ -124,9 +120,7 @@ def set_concurrency(self, new_concurrency: int) -> None: return if new_concurrency < 0 or new_concurrency > self.total_capacity: - raise ValueError( - f"Concurrency must be between 0 and {self.total_capacity}" - ) + raise ValueError(f"Concurrency must be between 0 and {self.total_capacity}") old_concurrency = self.concurrency old_buffer = self.buffer_concurrency @@ -423,7 +417,9 @@ async def _fetch_and_process_jobs(self): ) finally: if should_release: - if (not job or self._stop_event.is_set()) and acquire_sem_task.done(): + if ( + not job or self._stop_event.is_set() + ) and acquire_sem_task.done(): self._job_semaphore.release() self._new_job_event.clear()