diff --git a/tests/test_command.py b/tests/test_command.py index d808520..08fbed1 100644 --- a/tests/test_command.py +++ b/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( diff --git a/tests/test_executor.py b/tests/test_executor.py index b7af974..298a4df 100644 --- a/tests/test_executor.py +++ b/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.""" diff --git a/tests/test_signals.py b/tests/test_signals.py index e94b289..2a3ed9d 100644 --- a/tests/test_signals.py +++ b/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) diff --git a/threadmill/executor.py b/threadmill/executor.py index f9490b3..0f0bbb1 100644 --- a/threadmill/executor.py +++ b/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() diff --git a/threadmill/management/commands/threadmill.py b/threadmill/management/commands/threadmill.py index c1001d6..7c40219 100644 --- a/threadmill/management/commands/threadmill.py +++ b/threadmill/management/commands/threadmill.py @@ -1,3 +1,4 @@ +import logging import signal import sys @@ -10,7 +11,7 @@ task_backends, ) -from ...executor import TaskExecutor +from ...executor import JsonFormatter, TaskExecutor def kill_softly(signum, frame): @@ -71,6 +72,13 @@ def add_arguments(self, parser): action="store_true", help="Drain the task queue and exit with 0.", ) + parser.add_argument( + "--log-format", + help=( + "Logging format string for worker log records, e.g." + " '%%(levelname)s %%(message)s'. Defaults to JSON." + ), + ) def handle( self, @@ -83,6 +91,7 @@ def handle( max_tasks, max_tasks_jitter, exit_empty, + log_format, **options, ): match sys.platform: @@ -101,6 +110,17 @@ def handle( raise CommandError( f"Backend does not support all specified queues: {_non_queues!r}" ) + try: + log_formatter = ( + JsonFormatter() if log_format is None else logging.Formatter(log_format) + ) + log_formatter.format( + logging.LogRecord( + "threadmill", logging.INFO, __file__, 1, "Ready", (), None + ) + ) + except (TypeError, ValueError) as e: + raise CommandError(f"Invalid log format: {log_format!r}") from e exe = TaskExecutor( backend=backend, workers=workers, @@ -109,6 +129,7 @@ def handle( max_tasks_jitter=max_tasks_jitter, exit_empty=exit_empty, queues=queues, + log_formatter=log_formatter, ) try: exe.run()