Log worker records as JSON with a pluggable formatter

tests/test_command.py

@@ -1,4 +1,6 @@
 import argparse
+import logging
+import re
 import signal
 from unittest.mock import patch
 
@@ -7,9 +9,17 @@
 from django.tasks import default_task_backend
 
 from tests.testapp.tasks import compute_workload, io_workload, memory_workload
+from threadmill.executor import JsonFormatter, handler
 from threadmill.management.commands import threadmill
 
 
+@pytest.fixture(autouse=True)
+def restore_log_formatter():
+    """Restore the default JSON log formatter after each test."""
+    yield
+    handler.setFormatter(JsonFormatter())
+
+
 class TestKillSoftly:
     def test_kill_softly__raise_keyboard_interrupt_with_signal_name(self):
         """Raise KeyboardInterrupt with signal metadata in message."""
@@ -30,6 +40,65 @@ def test_add_arguments__register_all_worker_options(self):
         assert parsed_arguments.threads == 1
         assert parsed_arguments.max_tasks == 0
         assert parsed_arguments.max_tasks_jitter == 0
+        assert parsed_arguments.log_format is None
+
+    def test_call_command__log_format(self):
+        """Run the worker with the given log format string."""
+        call_command(
+            "threadmill",
+            "worker",
+            verbosity=0,
+            workers=1,
+            exit_empty=True,
+            log_format="%(levelname)s %(message)s",
+        )
+        record = logging.LogRecord(
+            "threadmill", logging.INFO, __file__, 1, "Hello %s", ("world",), None
+        )
+        assert handler.formatter.format(record) == "INFO Hello world"
+
+    @pytest.mark.parametrize("log_format", ["100% done", "%(message)s 100% done", 123])
+    def test_call_command__log_format__raise_command_error(self, log_format):
+        """Raise CommandError for a log format that fails to format a record."""
+        with pytest.raises(
+            CommandError, match=re.escape(f"Invalid log format: {log_format!r}")
+        ):
+            call_command(
+                "threadmill",
+                "worker",
+                verbosity=0,
+                workers=1,
+                exit_empty=True,
+                log_format=log_format,
+            )
+
+    def test_call_command__log_format__defaults_to_json(self):
+        """Default to the JSON log formatter."""
+        handler.setFormatter(logging.Formatter("%(message)s"))
+        call_command(
+            "threadmill",
+            "worker",
+            verbosity=0,
+            workers=1,
+            exit_empty=True,
+        )
+        assert isinstance(handler.formatter, JsonFormatter)
+
+    def test_call_command__log_format__empty_string(self):
+        """Honor an explicitly empty log format string."""
+        handler.setFormatter(JsonFormatter())
+        call_command(
+            "threadmill",
+            "worker",
+            verbosity=0,
+            workers=1,
+            exit_empty=True,
+            log_format="",
+        )
+        record = logging.LogRecord(
+            "threadmill", logging.INFO, __file__, 1, "Hello %s", ("world",), None
+        )
+        assert handler.formatter.format(record) == "Hello world"
 
     @pytest.mark.benchmark
     def test_call_command__benchmark_compute(

tests/test_executor.py

@@ -1,5 +1,6 @@
 import dataclasses
 import datetime
+import json
 import logging
 import multiprocessing
 import sys
@@ -27,7 +28,13 @@
     echo,
 )
 from threadmill.backends.base import Broker
-from threadmill.executor import TaskExecutor, WorkerProcess, WorkerThread
+from threadmill.executor import (
+    JsonFormatter,
+    TaskExecutor,
+    WorkerProcess,
+    WorkerThread,
+    handler,
+)
 
 
 @task(queue_name="default")
@@ -70,9 +77,80 @@ def _make_worker(*, max_tasks: int | None = None) -> WorkerProcess:
         max_tasks=max_tasks,
         backend_alias="default",
         queues=("default",),
+        log_formatter=JsonFormatter(),
     )
 
 
+class TestJsonFormatter:
+    """Tests for the JsonFormatter class."""
+
+    def test_format__returns_json_payload(self) -> None:
+        """Return a JSON object with structured record fields."""
+        record = logging.LogRecord(
+            "multiprocessing",
+            logging.INFO,
+            __file__,
+            1,
+            "Task successful %r",
+            ("abc",),
+            None,
+        )
+        payload = json.loads(JsonFormatter().format(record))
+        assert set(payload) == {
+            "created_at",
+            "level",
+            "logger",
+            "message",
+            "process",
+            "process_name",
+            "thread",
+        }
+        assert payload["message"] == "Task successful 'abc'"
+        assert payload["level"] == "INFO"
+        assert payload["logger"] == "multiprocessing"
+        assert payload["thread"] == "MainThread"
+        assert (
+            datetime.datetime.fromisoformat(payload["created_at"]).utcoffset()
+            == datetime.datetime.fromtimestamp(
+                record.created, tz=timezone.get_current_timezone()
+            ).utcoffset()
+        )
+
+    def test_format__includes_extra_attributes(self) -> None:
+        """Include extra record attributes in the JSON payload."""
+        record = logging.LogRecord(
+            "multiprocessing",
+            logging.INFO,
+            __file__,
+            1,
+            "Task successful %r",
+            ("abc",),
+            None,
+        )
+        record.request_id = "abc"
+        record.duration_ms = 42
+        payload = json.loads(JsonFormatter().format(record))
+        assert payload["request_id"] == "abc"
+        assert payload["duration_ms"] == 42
+
+    def test_format__includes_exception_traceback(self) -> None:
+        """Include the exception traceback when the record carries exc_info."""
+        try:
+            raise ValueError("boom")
+        except ValueError:
+            record = logging.LogRecord(
+                "multiprocessing",
+                logging.ERROR,
+                __file__,
+                1,
+                "Task failed %r",
+                ("abc",),
+                sys.exc_info(),
+            )
+        payload = json.loads(JsonFormatter().format(record))
+        assert "ValueError: boom" in payload["exception"]
+
+
 class TestTaskExecutor:
     """Tests for the TaskExecutor dataclass and its methods."""
 
@@ -132,6 +210,7 @@ def test_create_worker_process__starts_worker(self):
         executor = TaskExecutor(backend=default_task_backend, queues=("default",))
         worker = executor.create_worker_process()
         assert worker.is_alive()
+        assert worker.log_formatter is executor.log_formatter
         worker.shutdown()
 
     def test_run__processes_enqueued_tasks_end_to_end(self):
@@ -148,6 +227,7 @@ def test_run__processes_enqueued_tasks_end_to_end(self):
         run_thread = threading.Thread(target=executor.run, daemon=True)
         run_thread.start()
         time.sleep(2)
+        assert handler.formatter is executor.log_formatter
         executor.shutdown()
         run_thread.join(timeout=5)
         assert not run_thread.is_alive()
@@ -286,6 +366,16 @@ def test_shutdown_requested__is_settable(self):
         worker.shutdown_requested.set()
         assert worker.shutdown_requested.is_set()
 
+    def test_run__applies_log_formatter_and_stops(self):
+        """run() applies the log formatter and returns when shutdown is requested."""
+        worker = _make_worker()
+        worker.shutdown_requested.set()
+        run_thread = threading.Thread(target=worker.run)
+        run_thread.start()
+        run_thread.join(timeout=5)
+        assert not run_thread.is_alive()
+        assert handler.formatter is worker.log_formatter
+
 
 class TestWorkerThread:
     """Tests for the WorkerThread class."""

tests/test_signals.py

@@ -10,7 +10,12 @@
 from django.utils import timezone
 
 from tests.testapp.tasks import echo
-from threadmill.executor import TaskExecutor, WorkerProcess, WorkerThread
+from threadmill.executor import (
+    JsonFormatter,
+    TaskExecutor,
+    WorkerProcess,
+    WorkerThread,
+)
 
 
 def _task_result(task, *args) -> TaskResult:
@@ -75,7 +80,10 @@ def test_execute_task_result__closes_stale_connection(self, stale_connection):
             stale_connection, "close", wraps=stale_connection.close
         ) as spy:
             worker = WorkerProcess(
-                thread_count=1, backend_alias="default", queues=("default",)
+                thread_count=1,
+                backend_alias="default",
+                queues=("default",),
+                log_formatter=JsonFormatter(),
             )
             thread = WorkerThread(worker=worker, index=0, backend=None)
 

threadmill/executor.py

@@ -3,6 +3,7 @@
 import asyncio
 import dataclasses
 import datetime
+import json
 import logging
 import multiprocessing
 import random
@@ -16,6 +17,7 @@
 from traceback import format_exception
 
 import django
+from django.core.serializers.json import DjangoJSONEncoder
 from django.tasks import TaskResult, task_backends
 from django.tasks.base import TaskContext, TaskError, TaskResultStatus
 from django.tasks.signals import task_finished, task_started
@@ -25,12 +27,62 @@
 if typing.TYPE_CHECKING:
     from .backends.base import Broker, ThreadmillTaskBackend
 
+
+class JsonFormatter(logging.Formatter):
+    """Format log records as single-line JSON objects."""
+
+    standard_attributes = frozenset(
+        {
+            "args",
+            "asctime",
+            "created",
+            "exc_info",
+            "exc_text",
+            "filename",
+            "funcName",
+            "levelname",
+            "levelno",
+            "lineno",
+            "message",
+            "module",
+            "msecs",
+            "msg",
+            "name",
+            "pathname",
+            "process",
+            "processName",
+            "relativeCreated",
+            "stack_info",
+            "taskName",
+            "thread",
+            "threadName",
+        }
+    )
+
+    def format(self, record: logging.LogRecord) -> str:
+        payload = {
+            "created_at": datetime.datetime.fromtimestamp(
+                record.created, tz=timezone.get_current_timezone()
+            ),
+            "level": record.levelname,
+            "logger": record.name,
+            "message": record.getMessage(),
+            "process": record.process,
+            "process_name": record.processName,
+            "thread": record.threadName,
+        } | {
+            key: value
+            for key, value in record.__dict__.items()
+            if key not in self.standard_attributes
+        }
+        if record.exc_info:
+            payload["exception"] = "".join(format_exception(*record.exc_info))
+        return json.dumps(payload, cls=DjangoJSONEncoder)
+
+
 logger = multiprocessing.get_logger()
-formatter = logging.Formatter(
-    "%(levelname)s: %(asctime)s - pid=%(process)s - %(message)s"
-)
 handler = logging.StreamHandler()
-handler.setFormatter(formatter)
+handler.setFormatter(JsonFormatter())
 logger.addHandler(handler)
 logger.setLevel(logging.INFO)
 
@@ -53,6 +105,7 @@ class TaskExecutor:
     queues: tuple[str]
     broker: Broker | None = dataclasses.field(default=None, init=False)
     exit_empty: bool = False
+    log_formatter: logging.Formatter = dataclasses.field(default_factory=JsonFormatter)
 
     def __post_init__(self) -> None:
         """Initialize derived orchestration fields and queues."""
@@ -69,17 +122,19 @@ def get_maximum_tasks_per_child(self) -> int | None:
     def create_worker_process(self) -> WorkerProcess:
         """Create and start a new worker process."""
         worker = WorkerProcess(
-            self.thread_count,
-            self.get_maximum_tasks_per_child(),
-            self.backend.alias,
-            self.queues,
-            self.exit_empty,
+            thread_count=self.thread_count,
+            max_tasks=self.get_maximum_tasks_per_child(),
+            backend_alias=self.backend.alias,
+            queues=self.queues,
+            exit_empty=self.exit_empty,
+            log_formatter=self.log_formatter,
         )
         worker.start()
         return worker
 
     def run(self) -> None:
         """Start consuming tasks until shutdown is requested."""
+        handler.setFormatter(self.log_formatter)
         self.worker_processes = [
             self.create_worker_process() for _ in range(self.process_count)
         ]
@@ -127,11 +182,13 @@ class WorkerProcess(multiprocessing.Process):
 
     def __init__(
         self,
+        *,
         thread_count: int,
         max_tasks: int | None = None,
         backend_alias: str = "",
         queues: tuple[str, ...] = (),
         exit_empty: bool = False,
+        log_formatter: logging.Formatter,
     ) -> None:
         """Create process with dedicated thread pool for task execution."""
         self.shutdown_requested = multiprocessing.Event()
@@ -141,12 +198,14 @@ def __init__(
         self.backend_alias = backend_alias
         self.queues = queues
         self.exit_empty = exit_empty
+        self.log_formatter = log_formatter
         self.task_count = 0
         self.lock: threading.Lock | None = None
         self.expired: threading.Event | None = None
 
     def run(self) -> None:
         """Start consumer execution inside this process."""
+        handler.setFormatter(self.log_formatter)
         django.setup()
         logger.info("Starting worker process %s", self.name)
         self.lock = threading.Lock()