Minimal repro: hatchet-sdk Python worker log capture retains every log record for the worker process lifetime

hatchet_leak_repro.py

"""Minimal reproduction: hatchet-sdk's worker log capture retains every log
record for the lifetime of the worker process.

Run:  pip install hatchet-sdk && python hatchet_leak_repro.py

No Hatchet engine, token, or network access required. No user-defined logging
handler is involved: the only handler in play is the one `hatchet_sdk` installs
on its own inside `capture_logs`.

What this measures
------------------
`WorkerActionRunLoopManager.aio_start` wraps the worker's *entire* action run
loop in `capture_logs` (worker/runner/run_loop_manager.py):

    async def aio_start(self, retry_count: int = 1) -> None:
        if self.client.config.disable_log_capture:
            await self._async_start()
        else:
            await capture_logs(
                self.client.log_interceptor, self.log_sender, self._async_start
            )()

`capture_logs` creates one `StringIO`, wraps it in a `LogForwardingHandler`
(a `logging.StreamHandler`), attaches it to the configured logger, and closes
the stream only in its `finally` — i.e. when the worker stops. `emit()` calls
`super().emit(record)`, which appends the formatted record to that `StringIO`.
Nothing in the SDK ever reads, seeks, or truncates it.

So the workload below is exactly "a worker that logs while it runs", and the
growth is measured as process RSS, inside the run, where a live worker sits for
its whole lifetime.
"""

import asyncio
import gc
import logging
import resource
import sys
from io import StringIO

from hatchet_sdk.worker.runner.utils.capture_logs import (
    LogForwardingHandler,
    capture_logs,
)

LINE_BYTES = 2_400  # a structured JSON log line; ours average 2,277 bytes
COUNTS = (0, 10_000, 20_000, 40_000, 80_000)


class StubLogSender:
    """Stands in for AsyncLogSender.

    Deliberately inert: it drops everything. That removes the SDK's outbound log
    queue (which *is* bounded, `log_queue_size`, and drops on overflow) as an
    explanation, isolating the retention to the StringIO alone.
    """

    def publish(self, record: object) -> None:
        pass


def rss_bytes() -> int:
    """Current resident set size.

    On Linux this reads /proc, so it is current RSS rather than a high-water
    mark — the growth below is retained memory, not a transient peak.
    """
    try:
        with open("/proc/self/statm") as statm:
            return int(statm.read().split()[1]) * resource.getpagesize()
    except OSError:
        usage = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
        return usage if sys.platform == "darwin" else usage * 1024


async def measure(n_records: int, capture: bool) -> tuple[float, int]:
    """Log `n_records` lines the way a running worker does; return (RSS MiB
    retained, chars sitting in the SDK's buffer)."""
    logger = logging.getLogger(f"repro.{capture}.{n_records}")
    logger.handlers.clear()
    logger.setLevel(logging.INFO)
    logger.propagate = False
    line = "x" * LINE_BYTES
    result: dict[str, float | int] = {}

    async def worker_run_loop() -> None:
        # Stand-in for _async_start: what the worker does for its whole life.
        gc.collect()
        before = rss_bytes()
        for _ in range(n_records):
            logger.info(line)
        gc.collect()
        result["rss"] = (rss_bytes() - before) / 1024 / 1024
        # Read the buffer the SDK made for itself, before its finally closes it.
        buffered = 0
        for handler in logger.handlers:
            if isinstance(handler, LogForwardingHandler):
                handler.flush()  # a StreamHandler flush reclaims nothing here
                buffered = handler.stream.tell()
        result["buffered"] = buffered

    if capture:
        # Exactly what aio_start does when disable_log_capture is False.
        await capture_logs(logger, StubLogSender(), worker_run_loop)()
    else:
        # Exactly what aio_start does when disable_log_capture is True.
        await worker_run_loop()

    return float(result["rss"]), int(result["buffered"])


async def main() -> None:
    import importlib.metadata

    print(
        f"hatchet-sdk {importlib.metadata.version('hatchet-sdk')} "
        f"| python {sys.version.split()[0]} | {sys.platform}\n"
    )

    print(f"Each record is a {LINE_BYTES}-byte log line emitted on the logger the")
    print("SDK was configured with, from inside the wrapped run loop.\n")
    header = f"{'records':>9} | {'log text':>10} | {'SDK buffer':>14} | {'RSS kept':>10} | {'control':>8}"
    print(header)
    print("-" * len(header))

    for n in COUNTS:
        leaked_rss, buffered = await measure(n, capture=True)
        control_rss, control_buffered = await measure(n, capture=False)
        assert control_buffered == 0
        print(
            f"{n:>9,} | {n * LINE_BYTES / 1024 / 1024:>6.0f} MiB | "
            f"{buffered:>14,} | {leaked_rss:>6.1f} MiB | {control_rss:>4.1f} MiB"
        )

    print(
        "\n'control' is the same workload with capture_logs skipped, i.e. what"
        "\n`disable_log_capture=True` gives you. Retained RSS tracks log volume"
        "\nlinearly with capture on, and is flat with it off."
    )

    # The buffer is not incidental to an odd configuration: the SDK's default
    # logger is the root logger, so out of the box this collects every stdlib
    # record the process emits.
    from hatchet_sdk.config import ClientConfig

    default_logger = ClientConfig.model_fields["logger"].default
    print(
        f"\nClientConfig.logger default -> {default_logger!r} "
        f"(is root: {default_logger is logging.getLogger()})"
    )


if __name__ == "__main__":
    asyncio.run(main())