Skip to content
Open
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 15 additions & 3 deletions src/rhapsody/telemetry/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,9 +176,13 @@ def __init__(
# correct parent span context in the dispatch loop.
self._task_contexts: dict[str, Any] = {}

# Token returned by context.attach() when the session span is activated in start().
# Stored so it can be properly detached in stop().
# Token returned by context.attach() when the session span is activated in start(),
# plus the asyncio task that performed the attach. ContextVar tokens can only be
# reset from the contextvars.Context they were created in (i.e. the same task), so
# stop() skips the detach when it runs in a different task — OTel would otherwise
# log a spurious "Failed to detach context" ValueError at shutdown.
self._session_ctx_token: Any = None
self._session_ctx_task: Any = None

# Caller-supplied OTel extension points wired into RHAPSODY's internal providers
# at start() time alongside SpanBuffer / InMemoryMetricReader.
Expand Down Expand Up @@ -251,6 +255,7 @@ async def start(self) -> None:
self._session_ctx_token = otel_context.attach(
trace_mod.set_span_in_context(self._session_span)
)
self._session_ctx_task = asyncio.current_task()

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Storing a strong reference to an asyncio.Task object can prevent the task, its coroutine, and all local variables in its call stack from being garbage collected, potentially causing a memory leak in long-running applications. Using a weak reference instead avoids this issue while still allowing identity comparison.

Suggested change
self._session_ctx_task = asyncio.current_task()
import weakref
_task = asyncio.current_task()
self._session_ctx_task = weakref.ref(_task) if _task is not None else None

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 9161618 — task is now held as a weakref.ref, with the import at module level.


self._running = True
self._dispatch_task = asyncio.create_task(self._dispatch_loop(), name="telemetry-dispatch")
Expand Down Expand Up @@ -322,8 +327,15 @@ async def stop(self) -> None:
from opentelemetry import context as otel_context

if self._session_ctx_token is not None:
otel_context.detach(self._session_ctx_token)
# Detach only when stop() runs in the task that attached the token in
# start(); a token reset from any other task raises ValueError inside
# otel_context.detach(), which OTel catches and logs as an ERROR.
# Skipping is safe: the attach only ever affected the starting task's
# context, and the tracer provider shuts down right below.
if asyncio.current_task() is self._session_ctx_task:
otel_context.detach(self._session_ctx_token)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Dereference the weak reference to the task before performing the identity check. If the task has already been garbage collected, the weak reference will resolve to None, and the detach step will be safely skipped.

Suggested change
if asyncio.current_task() is self._session_ctx_task:
otel_context.detach(self._session_ctx_token)
_attached_task = self._session_ctx_task() if self._session_ctx_task is not None else None
if _attached_task is not None and asyncio.current_task() is _attached_task:
otel_context.detach(self._session_ctx_token)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 9161618 — the weakref is dereferenced before the identity check; a collected task resolves to None and the detach is skipped.

self._session_ctx_token = None
self._session_ctx_task = None

if self._meter_provider:
self._meter_provider.shutdown()
Expand Down
Loading