[BUG]: grpc aio client interceptor leaks unfinished spans on cancellation and on un-iterated streams

[jahodfra]

Tracer Version(s)

4.11.1

Python Version(s)

Python 3.12.9

Pip Version(s)

uv 0.8

Bug Report

Summary

ddtrace/contrib/internal/grpc/aio_client_interceptor.py creates a span per RPC but
fails to finish it on two paths. Unfinished spans are retained by SpanAggregator
for the lifetime of the process, so a long-running service leaks memory in
proportion to how often these paths are hit.

1. Cancelled unary RPC. _wrap_unary_response catches only aio.AioRpcError.
   An asyncio.CancelledError escapes before the done-callback is registered, so
   span.finish() never runs.
2. Streaming response never iterated. _wrap_stream_response is an async
   generator and registers the done-callback inside its body, so it only runs on
   the first __anext__. A caller that never iterates the response leaves the span
   unfinished.
   Path 1 is the serious one: in any ASGI service behind a load balancer, a client
   disconnect cancels the request task, which cancels the in-flight gRPC call.

Environment

• ddtrace 4.11.1
• grpcio 1.82.1
• Python 3.12.9, macOS (arm64); also expected on Linux — nothing here is platform-specific
• Reproduced with ddtrace.patch(grpc=True) and with ddtrace-run

Reproducer

ddtrace_grpc_aio_span_leak.py, attached. No .proto and no generated code —
request_deserializer/response_serializer of None makes grpc pass bytes
through. Only grpcio and ddtrace are required.

python ddtrace_grpc_aio_span_leak.py 100

Actual output

ddtrace 4.11.1, grpcio 1.82.1, python 3.12.9, 100 RPCs each
  case                                   leaked      per RPC  by service
  unary, completed        (control)           0 (+0.000/RPC)  -
  unary, cancelled        (BUG 1)           200 (+2.000/RPC)  grpc-aio-client/unfinished=100, grpc-aio-server/finished=100
  stream, consumed        (control)           0 (+0.000/RPC)  -
  stream, never iterated  (BUG 2)           100 (+1.000/RPC)  grpc-aio-client/unfinished=100
total spans the tracer will never finish or send: 300
Shutting down tracer with 304 spans. These spans will not be sent to Datadog: ...

Expected output

0 leaked on every row.
The count is read from tracer._span_aggregator._traces, and ddtrace's own
shutdown warning reports the same spans. The leak is stable — re-measured at 2s,
10s and 30s after the RPCs complete, with identical counts, so these are not spans
still in flight.
The deterministic signal is grpc-aio-client/unfinished == N on the two bug rows;
that is exactly one leaked span per RPC and it reproduces every run. Whether
finished grpc-aio-server spans are retained alongside them varies between runs
(it depends on whether the server handler got far enough to create its span before
the client gave up), which is why the per RPC figure on those rows moves between
1.000 and 2.000. Both are the same underlying leak — see "Amplification" below.

Root cause

Bug 1 — _wrap_unary_response, aio_client_interceptor.py:200

async def _wrap_unary_response(self, continuation, span):
    try:
        call = await continuation()
        code = await call.code()
        details = await call.details()
        _handle_add_callback(call, _done_callback_unary(span, code, details))
        return call
    except aio.AioRpcError as rpc_error:
        _handle_rpc_error(span, rpc_error)
        raise

await continuation() / await call.code() can raise asyncio.CancelledError,
which is not an AioRpcError and (since 3.8) not an Exception. It propagates out
before _handle_add_callback, so nothing ever finishes the span. Affects
intercept_unary_unary and intercept_stream_unary.

Bug 2 — _wrap_stream_response, aio_client_interceptor.py:175

async def _wrap_stream_response(self, call, span):
    try:
        _handle_add_callback(call, _done_callback_stream(span))   # only on first __anext__
        async for response in call:
            yield response

intercept_unary_stream / intercept_stream_stream return this generator without
starting it. Zero iterations means the callback is never registered. Consuming even
one item is enough to avoid the leak; consuming none is not.

Amplification: one unfinished span pins its whole trace

SpanAggregator.on_span_finish deletes a trace only once every span in it has
finished:

trace.num_finished += 1
is_trace_complete = trace.num_finished >= len(trace.spans)
if is_trace_complete:
    del self._traces[trace_id]

So the retained memory is not one span but the entire trace. That is what the
grpc-aio-server/finished=100 column above shows: those server spans finished
correctly, but they share a trace with the unfinished client span and are retained
with it. In a real service the same trace also holds the FastAPI/ASGI span, DB
spans, and so on — all pinned by a single cancelled gRPC call.

Suggested fix

Finish the span on cancellation, e.g. catch BaseException (or
asyncio.CancelledError explicitly) in _wrap_unary_response, and register the
stream done-callback in intercept_unary_stream/intercept_stream_stream before
returning the generator rather than inside it.

Unrelated bug noticed in the same file

aio_client_interceptor.py:186:

except StopAsyncIteration:
    # Callback will handle span finishing
    _handle_cancelled_error()
    raise

_handle_cancelled_error is async def _handle_cancelled_error(call, span). Called
with zero arguments and not awaited, this raises TypeError if the branch is ever
reached.
Also, the consumed-stream path logs Unable to parse async grpc string for status code and details. on every RPC (from utils._parse_rpc_repr_string in
_done_callback_stream) — visible in the reproducer output above.

Reproduction Code

#!/usr/bin/env python
"""Reproducer: ddtrace's grpc aio client interceptor leaks spans that are never finished.

Two paths, both in ddtrace/contrib/internal/grpc/aio_client_interceptor.py:

1. CANCELLED UNARY RPC. _wrap_unary_response catches only aio.AioRpcError:

       try:
           call = await continuation()
           code = await call.code()
           details = await call.details()
           _handle_add_callback(call, _done_callback_unary(span, code, details))
           return call
       except aio.AioRpcError as rpc_error:
           _handle_rpc_error(span, rpc_error)
           raise

   An asyncio.CancelledError -- what asyncio.wait_for, a Starlette client
   disconnect, or task-group teardown raises -- escapes before the done-callback
   is registered, so span.finish() never runs.

2. STREAMING RESPONSE NEVER ITERATED. _wrap_stream_response is an async
   generator, and the done-callback registration lives inside its body:

       async def _wrap_stream_response(self, call, span):
           try:
               _handle_add_callback(call, _done_callback_stream(span))   # runs on first __anext__
               async for response in call:
                   yield response

   The interceptor returns the un-started generator, so a caller that never
   iterates the response leaves the span unfinished. Consuming even one item is
   enough to avoid it; consuming zero is not.

Unfinished spans are retained by the tracer's SpanAggregator (_traces is keyed by
trace_id and only deleted once every span in the trace finishes), so they
accumulate for the lifetime of the process. ddtrace reports them itself at exit:

    Shutting down tracer with 201 spans. These spans will not be sent to Datadog: ...

Impact: a service whose RPCs are ever cancelled -- e.g. an ASGI app behind a load
balancer, where a client disconnect cancels the request task -- leaks one span per
cancelled RPC, unbounded.

No .proto and no generated code needed: request_deserializer/response_serializer
of None makes grpc pass bytes through. Requires only grpcio and ddtrace.

    python ddtrace_grpc_aio_span_leak.py [N]

Expected: 0 leaked spans on every row. Actual: N on the two rows below.
"""

import asyncio
import sys
from collections import Counter

import grpc

import ddtrace
from ddtrace.trace import tracer

METHOD = "/leak.Echo/Call"
LOCAL = grpc.LocalConnectionType.LOCAL_TCP


def open_spans():
    """Spans the tracer is still holding because they were never finished.

    Split by service and by whether the span itself was ever finished. A trace is
    evicted only once every span in it finishes, so one unfinished span pins the
    whole trace -- including sibling spans that finished correctly.
    """
    aggregator = tracer._span_aggregator
    counts = Counter()
    with aggregator._lock:
        for trace in aggregator._traces.values():
            for span in trace.spans:
                counts[f"{span.service}/{'unfinished' if span.duration is None else 'finished'}"] += 1
    return counts


async def unary_echo(request, context):
    await asyncio.sleep(0.2)  # long enough for the client to cancel mid-flight
    return request


async def stream_echo(request, context):
    yield request


def server_handlers(kind):
    handler = grpc.unary_unary_rpc_method_handler if kind == "unary" else grpc.unary_stream_rpc_method_handler
    behavior = unary_echo if kind == "unary" else stream_echo
    return (grpc.method_handlers_generic_handler("leak.Echo", {"Call": handler(behavior, None, None)}),)


async def run(label, kind, drive, n):
    server = grpc.aio.server(handlers=server_handlers(kind))
    port = server.add_secure_port("127.0.0.1:0", grpc.local_server_credentials(LOCAL))
    await server.start()

    # ddtrace wraps grpc.aio.secure_channel, so the interceptors are installed here.
    channel = grpc.aio.secure_channel(f"127.0.0.1:{port}", grpc.local_channel_credentials(LOCAL))
    stub = (channel.unary_unary if kind == "unary" else channel.unary_stream)(METHOD, _registered_method=False)

    await drive(stub)
    await asyncio.sleep(1.0)
    before = open_spans()
    for _ in range(n):
        await drive(stub)
    await asyncio.sleep(2.0)
    after = open_spans()

    await channel.close()
    await server.stop(0)
    leaked = after - before  # Counter subtraction keeps only positive counts
    total = sum(leaked.values())
    detail = ", ".join(f"{svc}={cnt}" for svc, cnt in sorted(leaked.items())) or "-"
    print(f"  {label:36} {total:>8} ({total / n:+.3f}/RPC)  {detail}")
    return total


async def unary_completed(stub):
    await stub(b"ping")


async def unary_cancelled(stub):
    try:
        await asyncio.wait_for(stub(b"ping"), timeout=0.02)
    except (asyncio.TimeoutError, asyncio.CancelledError, grpc.aio.AioRpcError):
        pass


async def stream_consumed(stub):
    async for _ in stub(b"ping"):
        pass


async def stream_abandoned(stub):
    stub(b"ping")  # response never iterated


async def main(n):
    ddtrace.patch(grpc=True)
    print(f"ddtrace {ddtrace.__version__}, grpcio {grpc.__version__}, python {sys.version.split()[0]}, {n} RPCs each\n")
    print(f"  {'case':36} {'leaked':>8} {'per RPC':>12}  by service")

    leaked = 0
    leaked += await run("unary, completed        (control)", "unary", unary_completed, n)
    leaked += await run("unary, cancelled        (BUG 1)", "unary", unary_cancelled, n)
    leaked += await run("stream, consumed        (control)", "stream", stream_consumed, n)
    leaked += await run("stream, never iterated  (BUG 2)", "stream", stream_abandoned, n)

    print(f"\ntotal spans the tracer will never finish or send: {leaked}")
    print("the shutdown warning below reports the same spans")


if __name__ == "__main__":
    asyncio.run(main(int(sys.argv[1]) if len(sys.argv) > 1 else 100))

\\\

### Error Logs

_No response_

### Libraries in Use

grpc

### Operating System

_No response_

[Bardakor]

Implemented in #19716.

I reproduced both lifecycle gaps on current main: timing out a unary RPC left its client span unfinished before callback registration, and returning a streaming wrapper without iterating it left the callback registration inside the lazy generator unexecuted.

The PR finishes the unary span while re-raising CancelledError, and registers unary-stream and stream-stream callbacks before returning their wrappers. The regressions cover the unary timeout plus never-iterated server-streaming and bidirectional-streaming calls. The complete gRPC asyncio suite passes on Python 3.9 and 3.11, and scripts/lint checks passes.