Refactoring the NNG support solution v1.5: Simplify store

Author: a-tedeev

taskiq/brokers/nng/__init__.py

@@ -1,24 +1,25 @@
-from hub import HubConfig, NNGHub
-from protocol import (
+"""NNG broker package for taskiq."""
+from .hub import HubConfig, NNGHub
+from .protocol import (
     ControlMessage,
     ControlResponse,
     MessageKind,
     TaskEnvelope,
     WorkerState,
     WorkerStatus,
 )
-from storage import QueueFullError, SQLiteJournal, StoreConfig
+from .storage import InMemoryStore, QueueFullError, StoreConfig
 
 __all__ = [
-    'HubConfig',
-    'NNGHub',
-    'ControlMessage',
-    'ControlResponse',
-    'MessageKind',
-    'TaskEnvelope',
-    'WorkerState',
-    'WorkerStatus',
-    'QueueFullError',
-    'SQLiteJournal',
-    'StoreConfig',
+    "HubConfig",
+    "NNGHub",
+    "ControlMessage",
+    "ControlResponse",
+    "MessageKind",
+    "TaskEnvelope",
+    "WorkerState",
+    "WorkerStatus",
+    "QueueFullError",
+    "InMemoryStore",
+    "StoreConfig",
 ]

taskiq/brokers/nng/broker.py

@@ -17,7 +17,7 @@
 from taskiq.acks import AckableMessage
 from taskiq.message import BrokerMessage
 
-from protocol import (
+from .protocol import (
     ControlMessage,
     ControlResponse,
     TaskEnvelope,

taskiq/brokers/nng/hub.py

@@ -3,12 +3,11 @@
 
 Run as a standalone process::
 
-    taskiq-nng-hub --control-addr ipc:///tmp/taskiq-nng.ipc \\
-                   --task-db /var/lib/taskiq/tasks.db
+    taskiq-nng-hub --control-addr ipc:///tmp/taskiq-nng.ipc
 
 Or embed it in an application for testing::
 
-    hub = NNGHub(HubConfig(control_addr="ipc:///tmp/h.ipc", task_db=":memory:"))
+    hub = NNGHub(HubConfig(control_addr="ipc:///tmp/h.ipc"))
     await hub.start()
     ...
     await hub.stop()
@@ -18,13 +17,11 @@
 import argparse
 import asyncio
 import base64
-import json
 import logging
 import os
 import signal
 import time
 import uuid
-from concurrent.futures import ThreadPoolExecutor
 from contextlib import suppress
 from dataclasses import dataclass, field
 from typing import Any
@@ -34,13 +31,13 @@
 except ImportError:
     pynng = None  # type: ignore[assignment]
 
-from protocol import (
+from .protocol import (
     ControlMessage,
     ControlResponse,
     TaskEnvelope,
     WorkerState,
 )
-from storage import QueueFullError, SQLiteJournal, StoreConfig
+from .storage import InMemoryStore, QueueFullError, StoreConfig
 
 logger = logging.getLogger(__name__)
 
@@ -50,7 +47,7 @@ class HubConfig:
     """Configuration for :class:`NNGHub`."""
 
     control_addr: str
-    task_db: str
+    task_db: str = ""  # kept for API compat; ignored by in-memory store
     max_pending: int = 10_000
     heartbeat_timeout: float = 15.0
     lease_timeout: float = 20.0
@@ -77,20 +74,18 @@ class NNGHub:
     independent ``nng_ctx`` contexts running concurrently.  Each context
     handles one request-reply at a time, so N workers can
     register/heartbeat/ack simultaneously without queuing behind each other.
-    This is the key fix over the single-context (serial) Rep0 in v2.
 
     **Data plane** — One ``Push0`` socket per registered worker, dialed to
     the worker's own ``Pull0`` listen address.  The hub explicitly targets
-    the least-loaded worker instead of relying on NNG round-robin, giving
-    us load-aware routing.
+    the least-loaded worker instead of relying on NNG round-robin.
 
-    **Persistence** — :class:`~taskiq.brokers.nng_storage.SQLiteJournal` in
-    WAL mode.  All storage calls are executed on a single-threaded
-    ``ThreadPoolExecutor`` so the asyncio event loop is never blocked and
-    SQLite write serialisation is guaranteed.
+    **State** — :class:`~taskiq.brokers.nng.storage.InMemoryStore`.  All
+    store operations are synchronous and execute directly on the asyncio event
+    loop without blocking (no I/O, no syscalls).
 
-    **Recovery** — On startup, tasks leased to workers that died during the
-    previous hub session are automatically requeued.
+    **Recovery** — On startup, any tasks that were leased before the hub last
+    stopped (within the same process lifetime) are automatically requeued by
+    :meth:`~InMemoryStore.recover_dead_workers`.
     """
 
     def __init__(self, config: HubConfig) -> None:
@@ -105,9 +100,8 @@ def __init__(self, config: HubConfig) -> None:
                 "Install it with: pip install taskiq[nng]"
             )
         self.config = config
-        self.store = SQLiteJournal(
+        self.store = InMemoryStore(
             StoreConfig(
-                path=config.task_db,
                 max_pending=config.max_pending,
                 lease_timeout=config.lease_timeout,
                 backoff_cap=config.backoff_cap,
@@ -117,16 +111,12 @@ def __init__(self, config: HubConfig) -> None:
         self._ctrl_sock: Any = None  # pynng.Rep0
         self._worker_push: dict[str, Any] = {}  # worker_id -> pynng.Push0
         self._tasks: list[asyncio.Task[None]] = []
-        # Single-threaded executor: serialises all SQLite calls on one OS thread.
-        self._db_exec = ThreadPoolExecutor(
-            max_workers=1, thread_name_prefix="nng-db"
-        )
 
     # ── lifecycle ─────────────────────────────────────────────────────────────
 
     async def start(self) -> None:
         """Start the hub: recover orphaned tasks, open sockets, spawn loops."""
-        await self._db(self.store.recover_dead_workers, self.config.heartbeat_timeout)
+        self.store.recover_dead_workers(self.config.heartbeat_timeout)
 
         self._ctrl_sock = pynng.Rep0(listen=self.config.control_addr)
         self._ctrl_sock.recv_timeout = self.config.recv_timeout_ms
@@ -142,11 +132,7 @@ async def start(self) -> None:
                     self._control_handler(ctx), name=f"hub-ctrl-{i}"
                 ),
             )
-        logger.info(
-            "NNG hub started on %s (db=%s)",
-            self.config.control_addr,
-            self.config.task_db,
-        )
+        logger.info("NNG hub started on %s", self.config.control_addr)
 
     async def stop(self) -> None:
         """Gracefully stop all hub loops and close sockets."""
@@ -163,17 +149,8 @@ async def stop(self) -> None:
         if self._ctrl_sock is not None:
             with suppress(Exception):
                 self._ctrl_sock.close()
-        self._db_exec.shutdown(wait=True)
         logger.info("NNG hub stopped")
 
-    # ── DB helper ─────────────────────────────────────────────────────────────
-
-    async def _db(self, fn: Any, *args: Any, **kwargs: Any) -> Any:
-        loop = asyncio.get_running_loop()
-        return await loop.run_in_executor(
-            self._db_exec, lambda: fn(*args, **kwargs)
-        )
-
     # ── control plane ─────────────────────────────────────────────────────────
 
     async def _control_handler(self, ctx: Any) -> None:
@@ -216,28 +193,26 @@ async def _handle(self, raw: bytes) -> ControlResponse:  # noqa: PLR0911, C901
             return await self._handle_register(msg.payload)
 
         if msg.kind == "heartbeat":
-            await self._db(self.store.heartbeat, msg.payload["worker_id"])
+            self.store.heartbeat(msg.payload["worker_id"])
             return ControlResponse(ok=True, payload={"ok": True})
 
         if msg.kind == "unregister":
             return await self._handle_unregister(msg.payload["worker_id"])
 
         if msg.kind == "drain":
-            await self._db(self.store.mark_draining, msg.payload["worker_id"])
+            self.store.mark_draining(msg.payload["worker_id"])
             return ControlResponse(ok=True, payload={"draining": True})
 
         if msg.kind == "ack":
-            ok = await self._db(
-                self.store.ack,
+            ok = self.store.ack(
                 msg.payload["task_id"],
                 msg.payload["worker_id"],
                 msg.payload["lease_id"],
             )
             return ControlResponse(ok=ok, payload={"acked": ok})
 
         if msg.kind == "nack":
-            ok = await self._db(
-                self.store.nack,
+            ok = self.store.nack(
                 msg.payload["task_id"],
                 msg.payload["worker_id"],
                 msg.payload["lease_id"],
@@ -246,26 +221,25 @@ async def _handle(self, raw: bytes) -> ControlResponse:  # noqa: PLR0911, C901
             return ControlResponse(ok=ok, payload={"nacked": ok})
 
         if msg.kind == "status":
-            task = await self._db(self.store.get_task, msg.payload["task_id"])
-            return ControlResponse(ok=bool(task), payload=dict(task) if task else {})
+            task = self.store.get_task(msg.payload["task_id"])
+            return ControlResponse(ok=bool(task), payload=task or {})
 
         if msg.kind == "stats":
-            s = await self._db(self.store.stats)
-            return ControlResponse(ok=True, payload=s)
+            return ControlResponse(ok=True, payload=self.store.stats())
 
         return ControlResponse(ok=False, error=f"unknown kind: {msg.kind!r}")
 
     async def _handle_submit(self, payload: dict[str, Any]) -> ControlResponse:
         envelope = TaskEnvelope(**payload)
         try:
-            await self._db(self.store.submit, envelope)
+            self.store.submit(envelope)
             return ControlResponse(ok=True, payload={"task_id": envelope.task_id})
         except QueueFullError:
             return ControlResponse(ok=False, error="queue full")
 
     async def _handle_register(self, payload: dict[str, Any]) -> ControlResponse:
         worker = WorkerState(**payload)
-        await self._db(self.store.register_worker, worker)
+        self.store.register_worker(worker)
         if worker.worker_id not in self._worker_push:
             try:
                 sock = pynng.Push0(dial=worker.task_addr)
@@ -279,7 +253,7 @@ async def _handle_register(self, payload: dict[str, Any]) -> ControlResponse:
         return ControlResponse(ok=True, payload={"registered": True})
 
     async def _handle_unregister(self, worker_id: str) -> ControlResponse:
-        await self._db(self.store.unregister_worker, worker_id)
+        self.store.unregister_worker(worker_id)
         sock = self._worker_push.pop(worker_id, None)
         if sock is not None:
             with suppress(Exception):
@@ -302,13 +276,12 @@ async def _dispatch_loop(self) -> None:
 
     async def _dispatch_once(self) -> bool:
         """Dispatch up to ``dispatch_batch`` due tasks to available workers."""
-        due = await self._db(self.store.due_tasks, self.config.dispatch_batch)
+        due = self.store.due_tasks(self.config.dispatch_batch)
         if not due:
             return False
         sent_any = False
         for row in due:
-            worker = await self._db(
-                self.store.choose_worker,
+            worker = self.store.choose_worker(
                 self.config.routing_policy,
                 heartbeat_timeout=self.config.heartbeat_timeout,
             )
@@ -319,8 +292,7 @@ async def _dispatch_once(self) -> bool:
             lease_id = uuid.uuid4().hex
             lease_until = time.time() + self.config.lease_timeout
 
-            if not await self._db(
-                self.store.mark_leased,
+            if not self.store.mark_leased(
                 row["task_id"], worker_id, lease_id, lease_until,
             ):
                 continue  # concurrent dispatch race; task already taken
@@ -331,19 +303,14 @@ async def _dispatch_once(self) -> bool:
                     "No push socket for worker %s, requeueing %s",
                     worker_id, row["task_id"],
                 )
-                await self._db(
-                    self.store.nack,
-                    row["task_id"], worker_id, lease_id, "no socket",
-                )
+                self.store.nack(row["task_id"], worker_id, lease_id, "no socket")
                 continue
 
-            # Include the hub-generated lease_id so the worker can ack with
-            # the exact token.  Omitting it was the core correctness bug in v2.
             envelope = TaskEnvelope(
                 task_id=row["task_id"],
                 task_name=row["task_name"],
                 payload_b64=base64.b64encode(row["payload"]).decode("ascii"),
-                labels=json.loads(row["labels_json"]),
+                labels=row["labels"],
                 lease_id=lease_id,
                 attempts=int(row["attempts"]) + 1,
                 max_retries=int(row["max_retries"]),
@@ -360,8 +327,7 @@ async def _dispatch_once(self) -> bool:
                     "Failed to deliver %s to worker %s: %s",
                     row["task_id"], worker_id, exc,
                 )
-                await self._db(
-                    self.store.nack,
+                self.store.nack(
                     row["task_id"], worker_id, lease_id,
                     f"dispatch send failed: {exc}",
                 )
@@ -373,11 +339,10 @@ async def _reaper_loop(self) -> None:
         while not self._stop.is_set():
             try:
                 await asyncio.sleep(self.config.reaper_interval)
-                reaped = await self._db(self.store.reap_expired_leases)
+                reaped = self.store.reap_expired_leases()
                 if reaped:
                     logger.debug("Reaped %d expired leases", reaped)
-                recovered = await self._db(
-                    self.store.recover_dead_workers,
+                recovered = self.store.recover_dead_workers(
                     self.config.heartbeat_timeout,
                 )
                 if recovered:
@@ -400,11 +365,6 @@ def _build_config() -> HubConfig:
         default=os.getenv("NNG_CONTROL_ADDR", "ipc:///tmp/taskiq-nng.ipc"),
         help="NNG address the hub listens on.  Env: NNG_CONTROL_ADDR",
     )
-    p.add_argument(
-        "--task-db",
-        default=os.getenv("NNG_TASK_DB", "/tmp/taskiq-nng-tasks.db"),  # noqa: S108
-        help="Path to the SQLite WAL task journal.  Env: NNG_TASK_DB",
-    )
     p.add_argument(
         "--max-pending",
         type=int,
@@ -445,7 +405,6 @@ def _build_config() -> HubConfig:
     )
     return HubConfig(
         control_addr=args.control_addr,
-        task_db=args.task_db,
         max_pending=args.max_pending,
         heartbeat_timeout=args.heartbeat_timeout,
         lease_timeout=args.lease_timeout,