fix(core): count execution capacity in permits so a hung job cannot kill the scheduler

chronis/core/base/scheduler.py

@@ -13,6 +13,7 @@
 )
 from chronis.core.execution.callbacks import OnFailureCallback, OnSuccessCallback
 from chronis.core.execution.coordinator import ExecutionCoordinator
+from chronis.core.execution.worker_pool import WorkerPool
 from chronis.core.jobs.definition import JobDefinition, JobInfo
 from chronis.core.schedulers.next_run_calculator import NextRunTimeCalculator
 from chronis.core.state import JobStatus
@@ -87,23 +88,19 @@ def __init__(
         self._init_execution_coordinator()
 
     def _init_execution_coordinator(self) -> None:
-        """Build a fresh thread pool and coordinator.
+        """Build a fresh worker pool and coordinator.
 
-        Called again on restart: a ThreadPoolExecutor cannot be reused after shutdown, so
+        Called again on restart: a pool that has been shut down stops accepting work, so
         start() must rebuild it. The registries are shared by reference, so registered
         functions and handlers survive.
         """
-        from concurrent.futures import ThreadPoolExecutor
-
-        self._executor = ThreadPoolExecutor(
-            max_workers=self.max_workers, thread_name_prefix="chronis-worker-"
-        )
+        self._workers = WorkerPool(max_workers=self.max_workers, logger=self.logger)
         self._executor_shutdown = False
 
         self._execution_coordinator = ExecutionCoordinator(
             storage=self.storage,
             lock=self.lock,
-            executor=self._executor,
+            workers=self._workers,
             function_registry=self._job_registry,
             failure_handler_registry=self._failure_handler_registry,
             success_handler_registry=self._success_handler_registry,

chronis/core/execution/coordinator.py

@@ -3,7 +3,6 @@
 import asyncio
 import threading
 from collections.abc import Callable, Generator
-from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 from dataclasses import dataclass
 from datetime import timedelta
@@ -20,6 +19,7 @@
 )
 from chronis.core.execution.callbacks import OnFailureCallback, OnSuccessCallback
 from chronis.core.execution.job_executor import JobExecutor
+from chronis.core.execution.worker_pool import Slot, WorkerPool
 from chronis.core.jobs.definition import JobInfo
 from chronis.core.schedulers.next_run_calculator import NextRunTimeCalculator
 from chronis.core.state import JobStatus
@@ -120,7 +120,7 @@ def __init__(
         self,
         storage: JobStorageAdapter,
         lock: LockAdapter | None,
-        executor: ThreadPoolExecutor,
+        workers: WorkerPool,
         function_registry: dict[str, Callable],
         failure_handler_registry: dict[str, OnFailureCallback],
         success_handler_registry: dict[str, OnSuccessCallback],
@@ -132,7 +132,7 @@ def __init__(
     ) -> None:
         self.storage = storage
         self.lock = lock
-        self.executor = executor
+        self.workers = workers
         self.function_registry = function_registry
         self.failure_handler_registry = failure_handler_registry
         self.success_handler_registry = success_handler_registry
@@ -149,7 +149,11 @@ def __init__(
 
     def try_execute(self, job_data: dict[str, Any], on_complete: Callable[[str], None]) -> bool:
         """
-        Try to execute a job with distributed lock.
+        Try to execute a job: reserve capacity, claim it, run it.
+
+        Capacity is reserved *before* the job is claimed. Claiming a job we have no room to
+        run would strand it in RUNNING until its liveness lease expired; leaving it
+        SCHEDULED lets the next poll — or a node with spare capacity — take it instead.
 
         Args:
             job_data: Job data from storage
@@ -164,27 +168,38 @@ def try_execute(self, job_data: dict[str, Any], on_complete: Callable[[str], Non
 
         job_logger = self.logger.with_context(job_id=job_id, job_name=job_name)
 
-        with self._acquire_lock_context(lock_key) as lock_acquired:
-            if not lock_acquired:
-                return False
+        slot = self.workers.reserve()
+        if slot is None:
+            return False
 
-            try:
-                claimed = self._try_claim_job_with_cas(job_id, job_data)
-
-                if claimed is None:
+        started = False
+        try:
+            with self._acquire_lock_context(lock_key) as lock_acquired:
+                if not lock_acquired:
                     return False
 
-                self._track_running(job_id)
                 try:
-                    self._trigger_execution(claimed, job_logger, on_complete)
-                except Exception:
-                    self._untrack_running(job_id)
-                    raise
-                return True
+                    claimed = self._try_claim_job_with_cas(job_id, job_data)
 
-            except Exception as e:
-                job_logger.error("Job execution failed", error=str(e), exc_info=True)
-                return False
+                    if claimed is None:
+                        return False
+
+                    self._track_running(job_id)
+                    try:
+                        self._trigger_execution(claimed, job_logger, on_complete, slot)
+                    except Exception:
+                        self._untrack_running(job_id)
+                        raise
+                    started = True
+                    return True
+
+                except Exception as e:
+                    job_logger.error("Job execution failed", error=str(e), exc_info=True)
+                    return False
+        finally:
+            # Every path that did not hand the slot to a running job gives it straight back.
+            if not started:
+                slot.release()
 
     def _track_running(self, job_id: str) -> None:
         """Mark a job as executing on this node so heartbeat() keeps it alive."""
@@ -248,43 +263,46 @@ def _trigger_execution(
         claimed: ClaimedJob,
         job_logger: ContextLogger,
         on_complete: Callable[[str], None],
+        slot: Slot,
     ) -> None:
         """
-        Trigger job execution with automatic rollback on failure.
+        Start the job on its slot, rolling the job back if it cannot be started at all.
 
-        Routes async functions to the shared event loop and sync functions
-        to the thread pool executor.
+        Async functions run on the shared event loop, sync functions on a worker thread.
+        Both hold the same kind of slot, so ``max_workers`` bounds them together.
 
         A sync job that times out cannot be killed — Python has no safe way to stop a
-        thread — so the timeout marks the job failed and abandons the thread. The job is
-        released from this node's bookkeeping at that moment rather than when the thread
-        eventually finishes: a thread that never finishes would otherwise hold its queue
-        slot forever and, because the queue rejects duplicates, block its own retry from
-        ever being picked up here again.
+        thread — so the timeout marks the job failed and abandons the thread. Everything the
+        job was holding is handed back at that moment rather than when the thread eventually
+        returns (it may never): its slot, its place in the queue, and its heartbeat. Waiting
+        for the thread instead would let a handful of hung jobs consume every slot and stop
+        the node executing anything at all.
         """
         job_id = claimed.job_id
         func_name = claimed.func_name
         timeout_seconds = claimed.timeout_seconds
 
-        def _on_done(_future: Any) -> None:
+        def _finish() -> None:
             self._untrack_running(job_id)
             on_complete(job_id)
 
+        timer: threading.Timer | None = None
+
         try:
             if self._job_executor.is_async(func_name):
                 loop = self._job_executor.ensure_async_loop()
                 future = asyncio.run_coroutine_threadsafe(
                     self._execute_async(claimed, job_logger), loop
                 )
-                future.add_done_callback(_on_done)
+
+                def _on_async_done(_future: Any) -> None:
+                    slot.release()
+                    _finish()
+
+                future.add_done_callback(_on_async_done)
             else:
                 gate = _OutcomeGate() if timeout_seconds else None
 
-                future = self.executor.submit(
-                    self._execute_in_background, claimed, job_logger, gate
-                )
-                future.add_done_callback(_on_done)
-
                 if timeout_seconds and gate is not None:
                     timeout_gate = gate
 
@@ -299,11 +317,11 @@ def _on_timeout() -> None:
                             timeout_seconds=timeout_seconds,
                         )
 
-                        # Hand the slot back now: the abandoned thread may never return, and
-                        # _on_done would never fire. Both calls are idempotent, so the late
-                        # _on_done (if it ever comes) is harmless.
-                        self._untrack_running(job_id)
-                        on_complete(job_id)
+                        # Give back the capacity now. The thread is written off; its slot,
+                        # queue place and heartbeat must not be. abandon() and _finish() are
+                        # idempotent, so the thread returning later is harmless.
+                        slot.abandon()
+                        _finish()
                         self._handle_job_failure(
                             claimed,
                             JobTimeoutError(f"Job '{job_id}' timed out after {timeout_seconds}s"),
@@ -314,24 +332,56 @@ def _on_timeout() -> None:
                     timer.daemon = True
                     timer.start()
 
-                    # Cancel timer when the future completes normally
-                    future.add_done_callback(lambda f: timer.cancel())
+                armed_timer = timer
+
+                def _run() -> None:
+                    try:
+                        self._execute_in_background(claimed, job_logger, gate)
+                    finally:
+                        if armed_timer is not None:
+                            armed_timer.cancel()
+                        _finish()
+
+                # Slot.run releases the slot when the thread returns — including a thread
+                # that was abandoned and came back later, where release() is a no-op.
+                slot.run(_run, name=f"chronis-job-{job_id}")
 
         except Exception as submit_error:
+            if timer is not None:
+                timer.cancel()  # the job never started; don't let its timeout fire
+
             job_logger.warning(
-                "Executor submit failed, rolling back to SCHEDULED",
+                "Could not start job, rolling back to SCHEDULED",
                 error=str(submit_error),
                 error_type=type(submit_error).__name__,
             )
-            try:
-                self._update_job_status(job_id, JobStatus.SCHEDULED)
-            except Exception as rollback_error:
-                job_logger.error(
-                    "Failed to rollback job status after submit failure",
-                    error=str(rollback_error),
-                )
+            self._rollback_unstarted_claim(claimed, job_logger)
             raise submit_error
 
+    def _rollback_unstarted_claim(self, claimed: ClaimedJob, job_logger: ContextLogger) -> None:
+        """Undo a claim for a job that never actually started.
+
+        The claim already moved next_run_time on to the *next* occurrence, so restoring only
+        the status would quietly drop the run we were about to perform: the job sits in
+        SCHEDULED with a future run time, invisible to both the poller and misfire detection,
+        and simply misses this firing. Put the scheduled time back so the job is still due
+        and the next poll picks it up.
+        """
+        updates: JobUpdateData = {
+            "status": JobStatus.SCHEDULED.value,
+            "updated_at": canonical_utc_now(),
+        }
+        if claimed.original_scheduled_time:
+            updates["next_run_time"] = claimed.original_scheduled_time
+
+        try:
+            self.storage.update_job(claimed.job_id, updates)
+        except Exception as rollback_error:
+            job_logger.error(
+                "Failed to rollback job after it could not be started",
+                error=str(rollback_error),
+            )
+
     def _execute_in_background(
         self,
         claimed: ClaimedJob,

chronis/core/execution/job_executor.py

@@ -4,6 +4,7 @@
 import inspect
 import threading
 from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 
 from chronis.core.common.exceptions import FunctionNotRegisteredError
@@ -22,13 +23,16 @@ def __init__(
         self,
         function_registry: dict[str, Callable],
         logger: ContextLogger,
+        async_io_workers: int = 4,
     ) -> None:
         self.function_registry = function_registry
         self.logger = logger
 
         self._async_loop: asyncio.AbstractEventLoop | None = None
         self._async_thread: threading.Thread | None = None
         self._async_lock = threading.Lock()
+        self._async_io_workers = async_io_workers
+        self._async_io_executor: ThreadPoolExecutor | None = None
 
     def is_async(self, func_name: str) -> bool:
         """Check if a registered function is async."""
@@ -85,11 +89,25 @@ async def execute_async(self, job_data: dict[str, Any], job_logger: ContextLogge
             await coro
 
     def ensure_async_loop(self) -> asyncio.AbstractEventLoop:
-        """Get or create the shared event loop for async jobs (thread-safe)."""
+        """Get or create the shared event loop for async jobs (thread-safe).
+
+        The loop gets an executor we own and name. An async job's post-execution work
+        (storage writes, callbacks) is blocking, so it is pushed off the loop with
+        run_in_executor — which, left to itself, would silently spin up asyncio's *default*
+        executor: an unnamed pool, sized by a rule of thumb, that nothing in Chronis
+        configures or shuts down.
+        """
         if self._async_loop is None or self._async_loop.is_closed():
             with self._async_lock:
                 if self._async_loop is None or self._async_loop.is_closed():
                     self._async_loop = asyncio.new_event_loop()
+
+                    self._async_io_executor = ThreadPoolExecutor(
+                        max_workers=self._async_io_workers,
+                        thread_name_prefix="chronis-async-io-",
+                    )
+                    self._async_loop.set_default_executor(self._async_io_executor)
+
                     self._async_thread = threading.Thread(
                         target=self._async_loop.run_forever,
                         daemon=True,
@@ -139,4 +157,8 @@ async def _drain() -> None:
         self._async_loop = None
         self._async_thread = None
 
+        if self._async_io_executor is not None:
+            self._async_io_executor.shutdown(wait=wait)
+            self._async_io_executor = None
+
         return drained