Address review feedback

Author: deepika-awasthi

custom_worker_tuner/README.md

@@ -5,7 +5,7 @@ This sample gates on a fake DB pool: the worker only polls for a new
 activity when the pool has a free connection.
 
 ## What this sample is
-downstream.py - A static-capacity counter. Pretends to be a DB pool. Two methods: increment() (claim a slot, returns False if full), decrement() (release)
+db_pool.py - A static-capacity counter. Pretends to be a DB pool. Two methods: increment() (claim a slot, returns False if full), decrement() (release)
 supplier.py - The custom slot supplier. On reserve_slot it polls downstream.increment() until it succeeds. On release_slot it calls downstream.decrement()
 shared.py - A RunBatch workflow that runs N do_work activities in parallel. The activity just sleeps
 worker.py - Wires Downstream + DownstreamAwareSupplier into a WorkerTuner

custom_worker_tuner/db_pool.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import logging
+import threading
+
+logger = logging.getLogger(__name__)
+
+
+class FakeDatabaseConnectionPool:
+    """Pretend connection pool with a fixed capacity, backed by a Semaphore."""
+
+    def __init__(self, allowed_connections: int, name: str = "db") -> None:
+        self.allowed_connections = allowed_connections
+        self.name = name
+        self._connection_pool = threading.BoundedSemaphore(allowed_connections)
+        logger.info(
+            "FakeDatabaseConnectionPool ready: name=%s allowed_connections=%d",
+            name,
+            allowed_connections,
+        )
+
+    def acquire(self, blocking: bool = True) -> bool:
+        """Claim a connection. When blocking, waits until one is free."""
+        return self._connection_pool.acquire(blocking=blocking)
+
+    def release(self) -> None:
+        """Return a connection to the pool."""
+        self._connection_pool.release()
+
+    @property
+    def in_use(self) -> int:
+        """Derived from the semaphore — single source of truth."""
+        return self.allowed_connections - self._connection_pool._value

custom_worker_tuner/downstream.py

@@ -12,7 +12,7 @@
     SlotReserveContext,
 )
 
-from custom_worker_tuner.downstream import Downstream
+from custom_worker_tuner.db_pool import FakeDatabaseConnectionPool
 
 logger = logging.getLogger(__name__)
 
@@ -27,44 +27,38 @@ def __init__(self, slot_id: int) -> None:
         self.slot_id = slot_id
 
 
-class DownstreamAwareSupplier(CustomSlotSupplier):
-    def __init__(self, downstream: Downstream, poll_interval_ms: int = 100) -> None:
-        self.downstream = downstream
-        self.poll_interval_ms = poll_interval_ms
-        logger.info(
-            "DownstreamAwareSupplier ready: downstream=%s poll_interval_ms=%d",
-            downstream.name,
-            poll_interval_ms,
-        )
+class PoolSlotSupplier(CustomSlotSupplier):
+    """Hands out slots only when the backing pool has a free connection."""
+
+    def __init__(self, connection_pool: FakeDatabaseConnectionPool) -> None:
+        self.connection_pool = connection_pool
+        logger.info("PoolSlotSupplier ready: connection_pool=%s", connection_pool.name)
 
     async def reserve_slot(self, ctx: SlotReserveContext) -> SlotPermit:
-        """block downstream until it has capacity to get incremented and then grant a slot."""
+        """Block until the pool has capacity, then grant a slot."""
+        await asyncio.to_thread(self.connection_pool.acquire)
         slot_id = next(_slot_id_gen)
-        while not self.downstream.increment():
-            await asyncio.sleep(self.poll_interval_ms / 1000.0)
         self._log("reserve", slot_id, "ready to poll")
         return _Permit(slot_id)
 
     def try_reserve_slot(self, ctx: SlotReserveContext) -> SlotPermit | None:
-        """Eager path: can i run this activity right now?"""
-        if self.downstream.increment():
+        """Eager path: try to claim a slot without blocking."""
+        if self.connection_pool.acquire(blocking=False):
             slot_id = next(_slot_id_gen)
             self._log("reserve", slot_id, "eager dispatch")
             return _Permit(slot_id)
         return None
 
     def mark_slot_used(self, ctx: SlotMarkUsedContext) -> None:
-        """A task arrived for a reserved slot"""
         slot_id = getattr(ctx.permit, "slot_id", "?")
         self._log("used", slot_id, "activity running")
 
     def release_slot(self, ctx: SlotReleaseContext) -> None:
-        """Return the slot to the downstream."""
         slot_id = getattr(ctx.permit, "slot_id", "?")
         detail = "no task arrived" if ctx.slot_info is None else "activity done"
-        self.downstream.decrement()
+        self.connection_pool.release()
         self._log("release", slot_id, detail)
 
     def _log(self, event: str, slot_id, note: str) -> None:
-        count = f"{self.downstream.currently_connected}/{self.downstream.allowed_connections}"
+        count = f"{self.connection_pool.in_use}/{self.connection_pool.allowed_connections}"
         logger.info(f"{event:<8}  {count:>5}  {note}")

custom_worker_tuner/supplier.py

@@ -7,13 +7,12 @@
 from temporalio.envconfig import ClientConfig
 from temporalio.worker import FixedSizeSlotSupplier, Worker, WorkerTuner
 
-from custom_worker_tuner.downstream import Downstream
+from custom_worker_tuner.db_pool import FakeDatabaseConnectionPool
 from custom_worker_tuner.shared import TASK_QUEUE, RunBatch, do_work
-from custom_worker_tuner.supplier import DownstreamAwareSupplier
+from custom_worker_tuner.supplier import PoolSlotSupplier
 
-CAPACITY = 10  # number of connections allowed at a time
-POLL_INTERVAL_MS = 500
-LOG_LEVEL = "INFO"  # flip to "DEBUG" to see every increment/decrement
+CAPACITY = 10  # number of pool connections (and concurrent activities)
+LOG_LEVEL = "INFO"
 
 
 async def main() -> None:
@@ -27,8 +26,8 @@ async def main() -> None:
     config.setdefault("target_host", "localhost:7233")
     client = await Client.connect(**config)
 
-    downstream = Downstream(allowed_connections=CAPACITY, name="db")
-    supplier = DownstreamAwareSupplier(downstream, poll_interval_ms=POLL_INTERVAL_MS)
+    pool = FakeDatabaseConnectionPool(allowed_connections=CAPACITY, name="db")
+    supplier = PoolSlotSupplier(pool)
     tuner = WorkerTuner.create_composite(
         workflow_supplier=FixedSizeSlotSupplier(100),
         activity_supplier=supplier,
@@ -44,7 +43,7 @@ async def main() -> None:
         tuner=tuner,
     )
 
-    print(f"\nworker started — capacity={CAPACITY}, poll={POLL_INTERVAL_MS}ms\n")
+    print(f"\nworker started — capacity={CAPACITY}\n")
     print("TIME          EVENT     COUNT   DETAIL")
     print("─" * 60)
     await worker.run()