Add queuespec for worker starting

Author: ryanhiebert

qio/__main__.py

@@ -1,7 +1,9 @@
 from contextlib import suppress
 from time import sleep as time_sleep
+from typing import Annotated
 
 from pika import ConnectionParameters
+from typer import Argument
 from typer import Typer
 
 from . import routine
@@ -10,6 +12,7 @@
 from .pika.broker import PikaBroker
 from .pika.transport import PikaTransport
 from .qio import Qio
+from .queuespec import QueueSpec
 from .sleep import sleep
 from .worker import Worker
 
@@ -92,14 +95,29 @@ def monitor(raw: bool = False):
 
 
 @app.command()
-def worker():
+def worker(
+    queuespec: Annotated[
+        QueueSpec,
+        Argument(
+            parser=QueueSpec.parse,
+            help="Queue configuration in format 'queue=concurrency'. "
+            "Examples: 'production=10', 'api,background=5'",
+            metavar="QUEUE[,QUEUE2,...]=CONCURRENCY",
+        ),
+    ],
+):
+    """Start a worker process for the specified queue and concurrency.
+
+    The worker will process invocations from the specified queue,
+    as many at a time as specified by the concurrency.
+    """
     connection_params = ConnectionParameters()
     qio = Qio(
         broker=PikaBroker(connection_params),
         transport=PikaTransport(connection_params),
         default_queue="qio",
     )
-    Worker(qio, queue="qio", concurrency=3)()
+    Worker(qio, queuespec)()
 
 
 @app.command()

qio/queuespec.md

@@ -0,0 +1,31 @@
+# QueueSpec
+
+The `QueueSpec` class represents a parsed queue configuration for workers.
+
+## Usage
+
+```python
+from qio.queuespec import QueueSpec
+
+# Single queue
+spec = QueueSpec.parse("production=10")
+print(spec.queues)      # ['production']
+print(spec.concurrency) # 10
+
+# Multiple queues
+spec = QueueSpec.parse("high,medium,low=5")
+print(spec.queues)      # ['high', 'medium', 'low']
+print(spec.concurrency) # 5
+```
+
+## Format
+
+Queue specifications use the format: `queue[,queue2,...]=concurrency`
+
+## Examples
+
+| Input               | Queues                      | Concurrency |
+|---------------------|-----------------------------|-------------|
+| `production=10`     | `['production']`            | `10`        |
+| `api,background=5`  | `['api', 'background']`     | `5`         |
+| `high,medium,low=2` | `['high', 'medium', 'low']` | `2`         |

qio/queuespec.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass
+from typing import Self
+
+
+@dataclass
+class QueueSpec:
+    """The parsed specification for how a worker will process tasks.
+
+    The queues are all the queues that are assigned to this amount of
+    concurrency. The exact implementation is broker-specific, but invocations
+    will often be roughly taken in round-robin order between the queues.
+    """
+
+    queues: list[str]
+    concurrency: int
+
+    @classmethod
+    def parse(cls, value: str) -> Self:
+        """Parse a queue spec string.
+
+        Examples:
+
+        | queuespec          | queues               | concurrency      |
+        +--------------------+----------------------+------------------+
+        | queue2=5           | ['queue2']           | 5                |
+        | queue1,queue2=10   | ['queue1', 'queue2'] | 10               |
+        """
+
+        if not value or value.strip() == "":
+            raise ValueError("Queue spec cannot be empty")
+
+        if "=" not in value:
+            raise ValueError(
+                f"Invalid queue spec. Be sure to include both the queue and capacity. "
+                f"got: '{value}'"
+            )
+
+        # Left split instead of right split to make = invalid for queue names
+        raw_queues, raw_concurrency = value.split("=", 1)
+
+        try:
+            concurrency = int(raw_concurrency)
+        except ValueError:
+            raise ValueError(
+                f"Concurrency must be a positive integer, "
+                f"got: '{raw_concurrency}' in '{value}'"
+            ) from None
+        if concurrency <= 0:
+            raise ValueError(
+                f"Concurrency must be a positive integer, "
+                f"got: '{raw_concurrency}' in '{value}'"
+            )
+
+        queues = [q.strip() for q in raw_queues.split(",") if q.strip()]
+
+        if not queues:
+            raise ValueError(f"No valid queue names found in '{value}'")
+
+        return cls(queues=queues, concurrency=concurrency)

qio/queuespec_test.py

@@ -0,0 +1,67 @@
+import pytest
+
+from .queuespec import QueueSpec
+
+
+def test_single_queue_parsing():
+    spec = QueueSpec.parse("production=10")
+    assert spec.queues == ["production"]
+    assert spec.concurrency == 10
+
+
+def test_multi_queue_parsing():
+    spec = QueueSpec.parse("high,medium,low=5")
+    assert spec.queues == ["high", "medium", "low"]
+    assert spec.concurrency == 5
+
+
+def test_whitespace_handling():
+    spec = QueueSpec.parse("  api , background  = 8 ")
+    assert spec.queues == ["api", "background"]
+    assert spec.concurrency == 8
+
+
+def test_empty_spec_error():
+    with pytest.raises(ValueError, match="Queue spec cannot be empty"):
+        QueueSpec.parse("")
+
+
+def test_missing_equals_error():
+    with pytest.raises(ValueError, match="Invalid queue spec"):
+        QueueSpec.parse("production")
+
+
+def test_empty_concurrency_error():
+    with pytest.raises(ValueError, match="Concurrency must be a positive integer"):
+        QueueSpec.parse("queue=")
+
+
+def test_invalid_concurrency_error():
+    with pytest.raises(ValueError, match="Concurrency must be a positive integer"):
+        QueueSpec.parse("queue=abc")
+
+
+def test_zero_concurrency_error():
+    with pytest.raises(ValueError, match="Concurrency must be a positive integer"):
+        QueueSpec.parse("queue=0")
+
+
+def test_negative_concurrency_error():
+    with pytest.raises(ValueError, match="Concurrency must be a positive integer"):
+        QueueSpec.parse("queue=-5")
+
+
+def test_empty_queue_names_error():
+    with pytest.raises(ValueError, match="No valid queue names found"):
+        QueueSpec.parse("=10")
+
+
+def test_comma_only_error():
+    with pytest.raises(ValueError, match="Invalid queue spec"):
+        QueueSpec.parse(",")
+
+
+def test_empty_queue_names_in_list():
+    spec = QueueSpec.parse("queue1,,queue2=5")
+    assert spec.queues == ["queue1", "queue2"]
+    assert spec.concurrency == 5

qio/worker.md

@@ -0,0 +1,42 @@
+# Worker
+
+The `Worker` class processes invocations from message queues with specified concurrency.
+
+## Usage
+
+```python
+from qio.qio import Qio
+from qio.queuespec import QueueSpec
+from qio.worker import Worker
+
+# Create queue specification
+queuespec = QueueSpec.parse("production=5")
+
+# Create worker
+qio = Qio(broker=broker, transport=transport)
+worker = Worker(qio, queuespec)
+
+# Run worker (blocks until stopped)
+worker()
+```
+
+## Concurrency
+
+The worker spawns multiple runner threads based on the concurrency setting:
+- Each thread processes invocations independently
+- Prefetch limit matches concurrency to prevent overloading
+- Threads coordinate through internal task queues
+
+## Lifecycle
+
+1. **Initialization**: Validates single queue, creates consumer and threads
+2. **Running**: Starts all threads and monitors for completion
+3. **Shutdown**: Stops gracefully on `KeyboardInterrupt` or thread failure
+
+## CLI Integration
+
+```bash
+python -m qio worker production=10
+```
+
+The CLI uses `QueueSpec.parse()` to convert string arguments into `QueueSpec` objects.

qio/worker.py

@@ -13,20 +13,29 @@
 from .invocation import Invocation
 from .invocation import LocalInvocationSuspended
 from .qio import Qio
+from .queuespec import QueueSpec
 from .thread import Thread
 
 
 class Worker:
-    def __init__(self, qio: Qio, *, queue: str, concurrency: int):
+    def __init__(self, qio: Qio, queuespec: QueueSpec):
         self.__qio = qio
+
+        if not queuespec.queues:
+            raise ValueError("No queues specified in queuespec")
+        if len(queuespec.queues) != 1:
+            raise ValueError("Only one queue is supported")
+
         self.__tasks = Queue[Invocation | SendContinuation | ThrowContinuation]()
-        self.__consumer = self.__qio.consume(queue=queue, prefetch=concurrency)
+        self.__consumer = self.__qio.consume(
+            queue=queuespec.queues[0], prefetch=queuespec.concurrency
+        )
         self.__continuer_events = self.__qio.subscribe({LocalInvocationSuspended})
 
         # Start threads event queues are created
         self.__runner_threads = [
             Thread(target=self.__runner, name=f"qio-runner-{i + 1}")
-            for i in range(concurrency)
+            for i in range(queuespec.concurrency)
         ]
         self.__continuer_thread = Thread(target=self.__continuer, name="qio-continuer")
         self.__receiver_thread = Thread(target=self.__receiver, name="qio-receiver")

qio_test.py

@@ -29,7 +29,7 @@ def test_integration():
 
         # 1. Start worker process in the background
         worker = subprocess.Popen(
-            [sys.executable, "-m", "qio", "worker"],
+            [sys.executable, "-m", "qio", "worker", "qio=3"],
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
         )