Supports configuring and running a worker from the CLI (#22)

* Supports configuring and running a worker from the CLI

This ended up making me confront the fact that we didn't really deal
with how the worker knows which tasks to run.  From the client side it's
really easy, since you can just put something blindly on the docket and
we'll register it for you if it's a function.  From the worker side,
there's nowhere to really start from.

I don't think this is the final form of this, but it's a baby step.
This version supports a string path pointing to any iterable of
functions you set up, so you can have:

```
my_tasks = [task_a, task_b, task_c]
```

and then register them with `--tasks my.module:my_tasks`.  We should
really think a lot more about this.

Closes #8

* Updating the chaos tests

* Use the right redis

* Making the CLI utility a code utility

Author: chrisguidry

chaos/driver.py

@@ -49,18 +49,15 @@ async def main(
     workers: int = 4,
 ):
     with RedisContainer("redis:7.4.2") as redis_server:
+        redis_url = f"redis://{redis_server.get_container_host_ip()}:{redis_server.get_exposed_port(6379)}"
         docket = Docket(
             name=f"test-docket-{uuid4()}",
-            host=redis_server.get_container_host_ip(),
-            port=redis_server.get_exposed_port(6379),
-            db=0,
+            url=redis_url,
         )
         environment = {
             **os.environ,
-            "CHAOS_DOCKET_NAME": docket.name,
-            "CHAOS_REDIS_HOST": docket.host,
-            "CHAOS_REDIS_PORT": str(docket.port),
-            "CHAOS_REDIS_DB": str(docket.db),
+            "DOCKET_NAME": docket.name,
+            "DOCKET_URL": redis_url,
         }
 
         if tasks % producers != 0:
@@ -93,8 +90,19 @@ async def spawn_worker() -> Process:
             return await asyncio.create_subprocess_exec(
                 *python_entrypoint(),
                 "-m",
-                "chaos.worker",
-                env=environment | {"OTEL_SERVICE_NAME": "chaos-worker"},
+                "docket",
+                "worker",
+                "--docket",
+                docket.name,
+                "--url",
+                redis_url,
+                "--tasks",
+                "chaos.tasks:chaos_tasks",
+                env=environment
+                | {
+                    "OTEL_SERVICE_NAME": "chaos-worker",
+                    "DOCKET_WORKER_REDELIVERY_TIMEOUT": "5s",
+                },
                 stdout=subprocess.DEVNULL,
                 stderr=subprocess.DEVNULL,
             )

chaos/producer.py

@@ -16,10 +16,8 @@
 
 async def main(tasks_to_produce: int):
     docket = Docket(
-        name=os.environ["CHAOS_DOCKET_NAME"],
-        host=os.environ["CHAOS_REDIS_HOST"],
-        port=int(os.environ["CHAOS_REDIS_PORT"]),
-        db=int(os.environ["CHAOS_REDIS_DB"]),
+        name=os.environ["DOCKET_NAME"],
+        url=os.environ["DOCKET_URL"],
     )
     tasks_sent = 0
     while tasks_sent < tasks_to_produce:

chaos/tasks.py

@@ -20,3 +20,6 @@ async def hello(
 
 async def toxic():
     sys.exit(42)
+
+
+chaos_tasks = [hello, toxic]

chaos/worker.py

@@ -1,6 +1,15 @@
+import asyncio
+import enum
+import logging
+import socket
+from datetime import timedelta
+from typing import Annotated
+
 import typer
 
-from docket import __version__
+from . import __version__, tasks
+from .docket import Docket
+from .worker import Worker
 
 app: typer.Typer = typer.Typer(
     help="Docket - A distributed background task system for Python functions",
@@ -9,11 +18,172 @@
 )
 
 
+class LogLevel(enum.StrEnum):
+    DEBUG = "DEBUG"
+    INFO = "INFO"
+    WARNING = "WARNING"
+    ERROR = "ERROR"
+    CRITICAL = "CRITICAL"
+
+
+def duration(duration_str: str | timedelta) -> timedelta:
+    """
+    Parse a duration string into a timedelta.
+
+    Supported formats:
+    - 123 = 123 seconds
+    - 123s = 123 seconds
+    - 123m = 123 minutes
+    - 123h = 123 hours
+    - 00:00 = mm:ss
+    - 00:00:00 = hh:mm:ss
+    """
+    if isinstance(duration_str, timedelta):
+        return duration_str
+
+    if ":" in duration_str:
+        parts = duration_str.split(":")
+        if len(parts) == 2:  # mm:ss
+            minutes, seconds = map(int, parts)
+            return timedelta(minutes=minutes, seconds=seconds)
+        elif len(parts) == 3:  # hh:mm:ss
+            hours, minutes, seconds = map(int, parts)
+            return timedelta(hours=hours, minutes=minutes, seconds=seconds)
+        else:
+            raise ValueError(f"Invalid duration string: {duration_str}")
+    elif duration_str.endswith("s"):
+        return timedelta(seconds=int(duration_str[:-1]))
+    elif duration_str.endswith("m"):
+        return timedelta(minutes=int(duration_str[:-1]))
+    elif duration_str.endswith("h"):
+        return timedelta(hours=int(duration_str[:-1]))
+    else:
+        return timedelta(seconds=int(duration_str))
+
+
 @app.command(
     help="Start a worker to process tasks",
 )
-def worker() -> None:
-    print("TODO: Configure and start a worker")
+def worker(
+    tasks: Annotated[
+        list[str],
+        typer.Option(
+            "--tasks",
+            help=(
+                "The dotted path of a task collection to register with the docket. "
+                "This can be specified multiple times.  A task collection is any "
+                "iterable of async functions."
+            ),
+        ),
+    ] = ["docket.tasks:standard_tasks"],
+    docket_: Annotated[
+        str,
+        typer.Option(
+            "--docket",
+            help="The name of the docket",
+            envvar="DOCKET_NAME",
+        ),
+    ] = "docket",
+    url: Annotated[
+        str,
+        typer.Option(
+            help="The URL of the Redis server",
+            envvar="DOCKET_URL",
+        ),
+    ] = "redis://localhost:6379/0",
+    name: Annotated[
+        str | None,
+        typer.Option(
+            help="The name of the worker",
+            envvar="DOCKET_WORKER_NAME",
+        ),
+    ] = socket.gethostname(),
+    logging_level: Annotated[
+        LogLevel,
+        typer.Option(
+            help="The logging level",
+            envvar="DOCKET_LOGGING_LEVEL",
+        ),
+    ] = LogLevel.INFO,
+    prefetch_count: Annotated[
+        int,
+        typer.Option(
+            help="The number of tasks to request from the docket at a time",
+            envvar="DOCKET_WORKER_PREFETCH_COUNT",
+        ),
+    ] = 10,
+    redelivery_timeout: Annotated[
+        timedelta,
+        typer.Option(
+            parser=duration,
+            help="How long to wait before redelivering a task to another worker",
+            envvar="DOCKET_WORKER_REDELIVERY_TIMEOUT",
+        ),
+    ] = timedelta(minutes=5),
+    reconnection_delay: Annotated[
+        timedelta,
+        typer.Option(
+            parser=duration,
+            help=(
+                "How long to wait before reconnecting to the Redis server after "
+                "a connection error"
+            ),
+            envvar="DOCKET_WORKER_RECONNECTION_DELAY",
+        ),
+    ] = timedelta(seconds=5),
+    until_finished: Annotated[
+        bool,
+        typer.Option(
+            "--until-finished",
+            help="Exit after the current docket is finished",
+        ),
+    ] = False,
+) -> None:
+    logging.basicConfig(level=logging_level)
+    asyncio.run(
+        Worker.run(
+            docket_name=docket_,
+            url=url,
+            name=name,
+            prefetch_count=prefetch_count,
+            redelivery_timeout=redelivery_timeout,
+            reconnection_delay=reconnection_delay,
+            until_finished=until_finished,
+            tasks=tasks,
+        )
+    )
+
+
+@app.command(help="Adds a trace task to the Docket")
+def trace(
+    docket_: Annotated[
+        str,
+        typer.Option(
+            "--docket",
+            help="The name of the docket",
+            envvar="DOCKET_NAME",
+        ),
+    ] = "docket",
+    url: Annotated[
+        str,
+        typer.Option(
+            help="The URL of the Redis server",
+            envvar="DOCKET_URL",
+        ),
+    ] = "redis://localhost:6379/0",
+    message: Annotated[
+        str,
+        typer.Argument(
+            help="The message to print",
+        ),
+    ] = "Howdy!",
+) -> None:
+    async def run() -> None:
+        async with Docket(name=docket_, url=url) as docket:
+            execution = await docket.add(tasks.trace)(message)
+            print(f"Added trace task {execution.key!r} to the docket {docket.name!r}")
+
+    asyncio.run(run())
 
 
 @app.command(

src/docket/cli.py

@@ -1,3 +1,4 @@
+import importlib
 from contextlib import asynccontextmanager
 from datetime import datetime, timezone
 from types import TracebackType
@@ -6,6 +7,7 @@
     AsyncGenerator,
     Awaitable,
     Callable,
+    Iterable,
     ParamSpec,
     Self,
     TypeVar,
@@ -31,26 +33,35 @@
 P = ParamSpec("P")
 R = TypeVar("R")
 
+TaskCollection = Iterable[Callable[..., Awaitable[Any]]]
+
 
 class Docket:
     tasks: dict[str, Callable[..., Awaitable[Any]]]
 
     def __init__(
         self,
         name: str = "docket",
-        host: str = "localhost",
-        port: int = 6379,
-        db: int = 0,
-        password: str | None = None,
+        url: str = "redis://localhost:6379/0",
     ) -> None:
+        """
+        Args:
+            name: The name of the docket.
+            url: The URL of the Redis server.  For example:
+                - "redis://localhost:6379/0"
+                - "redis://user:password@localhost:6379/0"
+                - "redis://user:password@localhost:6379/0?ssl=true"
+                - "rediss://localhost:6379/0"
+                - "unix:///path/to/redis.sock"
+        """
         self.name = name
-        self.host = host
-        self.port = port
-        self.db = db
-        self.password = password
+        self.url = url
 
     async def __aenter__(self) -> Self:
-        self.tasks = {}
+        from .tasks import standard_tasks
+
+        self.tasks = {fn.__name__: fn for fn in standard_tasks}
+
         return self
 
     async def __aexit__(
@@ -63,12 +74,7 @@ async def __aexit__(
 
     @asynccontextmanager
     async def redis(self) -> AsyncGenerator[Redis, None]:
-        async with Redis(
-            host=self.host,
-            port=self.port,
-            db=self.db,
-            password=self.password,
-        ) as redis:
+        async with Redis.from_url(self.url) as redis:
             yield redis
 
     def register(self, function: Callable[..., Awaitable[Any]]) -> None:
@@ -78,6 +84,19 @@ def register(self, function: Callable[..., Awaitable[Any]]) -> None:
 
         self.tasks[function.__name__] = function
 
+    def register_collection(self, collection_path: str) -> None:
+        """
+        Register a collection of tasks.
+
+        Args:
+            collection_path: A path in the format "module:collection".
+        """
+        module_name, _, member_name = collection_path.rpartition(":")
+        module = importlib.import_module(module_name)
+        collection = getattr(module, member_name)
+        for function in collection:
+            self.register(function)
+
     @overload
     def add(
         self,