From 3c23ea33324d965a054002c7b96fb4e650329850 Mon Sep 17 00:00:00 2001 From: Alex Alderman Webb Date: Thu, 23 Jul 2026 18:07:27 +0200 Subject: [PATCH 1/2] ref(o11y): Support `sentry-sdk` configured with the streaming trace lifecycle (#757) Use the `sentry_sdk.traces` API when the streaming lifecycle is enabled. The previous API is a no-op when the lifecycle is enabled. --- clients/python/pyproject.toml | 2 +- .../python/src/taskbroker_client/registry.py | 17 ++-- clients/python/src/taskbroker_client/sdk.py | 82 +++++++++++++++++++ .../taskbroker_client/worker/workerchild.py | 58 +++++++------ uv.lock | 6 +- 5 files changed, 127 insertions(+), 38 deletions(-) create mode 100644 clients/python/src/taskbroker_client/sdk.py diff --git a/clients/python/pyproject.toml b/clients/python/pyproject.toml index 4b48b4f0..a53555ee 100644 --- a/clients/python/pyproject.toml +++ b/clients/python/pyproject.toml @@ -5,7 +5,7 @@ description = "Taskbroker python client and worker runtime" readme = "README.md" dependencies = [ "sentry-arroyo>=2.41.0", - "sentry-sdk[http2]>=2.43.0", + "sentry-sdk[http2]>=2.66.1", "sentry-protos>=0.26.1", "confluent_kafka>=2.3.0", "cronsim>=2.6", diff --git a/clients/python/src/taskbroker_client/registry.py b/clients/python/src/taskbroker_client/registry.py index 63e8ae4e..48e81536 100644 --- a/clients/python/src/taskbroker_client/registry.py +++ b/clients/python/src/taskbroker_client/registry.py @@ -6,7 +6,6 @@ from concurrent import futures from typing import Any, cast -import sentry_sdk from arroyo.backends.kafka import KafkaPayload from arroyo.types import BrokerValue, Topic from sentry_protos.taskbroker.v1.taskbroker_pb2 import TaskActivation @@ -20,6 +19,7 @@ from taskbroker_client.metrics import MetricsBackend from taskbroker_client.retry import Retry from taskbroker_client.router import TaskRouter +from taskbroker_client.sdk import start_span from taskbroker_client.task import ExternalTask, P, R, Task from taskbroker_client.types import ContextHook, ProducerFactory, ProducerProtocol @@ -172,15 +172,16 @@ def send_task( ) -> ProducerFuture: topic = self.topic - with sentry_sdk.start_span( - op=OP.QUEUE_PUBLISH, + with start_span( name=activation.taskname, + op=OP.QUEUE_PUBLISH, origin="taskworker", - ) as span: - span.set_data(SPANDATA.MESSAGING_DESTINATION_NAME, activation.namespace) - span.set_data(SPANDATA.MESSAGING_MESSAGE_ID, activation.id) - span.set_data(SPANDATA.MESSAGING_SYSTEM, "taskworker") - + attributes={ + SPANDATA.MESSAGING_DESTINATION_NAME: activation.namespace, + SPANDATA.MESSAGING_MESSAGE_ID: activation.id, + SPANDATA.MESSAGING_SYSTEM: "taskworker", + }, + ): produce_future = self._producer(topic).produce( Topic(name=topic), KafkaPayload(key=None, value=activation.SerializeToString(), headers=[]), diff --git a/clients/python/src/taskbroker_client/sdk.py b/clients/python/src/taskbroker_client/sdk.py new file mode 100644 index 00000000..609360e2 --- /dev/null +++ b/clients/python/src/taskbroker_client/sdk.py @@ -0,0 +1,82 @@ +from contextlib import nullcontext +from typing import Any, ContextManager + +import sentry_sdk +from sentry_sdk.scope import Scope +from sentry_sdk.traces import StreamedSpan +from sentry_sdk.tracing import NoOpSpan, Span, Transaction +from sentry_sdk.tracing_utils import has_span_streaming_enabled + + +def start_transaction( + name: str, + op: str, + origin: str, + attributes: dict[str, Any], + headers: dict[str, Any], + sampling_context: dict[str, Any], +) -> Transaction | NoOpSpan | StreamedSpan | ContextManager[Any]: + """Start a transaction, or a span if span streaming is enabled.""" + span = None + try: + is_span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) + if is_span_streaming: + sentry_sdk.traces.continue_trace(headers) + Scope.set_custom_sampling_context(sampling_context) + + return sentry_sdk.traces.start_span( + name=name, + attributes={ + "sentry.op": op, + "sentry.origin": origin, + **attributes, + }, + ) + + transaction = sentry_sdk.continue_trace( + environ_or_headers=headers, + op=op, + name=name, + origin=origin, + ) + + span = sentry_sdk.start_transaction(transaction, custom_sampling_context=sampling_context) + for key, value in attributes.items(): + span.set_data(key, value) + except Exception: + pass + + if span is None: + return nullcontext() + return span + + +def start_span( + name: str, op: str, origin: str, attributes: dict[str, Any] +) -> Span | StreamedSpan | ContextManager[Any]: + """Start a span in the currently active trace lifecycle.""" + try: + is_span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) + if is_span_streaming: + return sentry_sdk.traces.start_span( + name=name, + attributes={ + "sentry.op": op, + "sentry.origin": origin, + **attributes, + }, + ) + + span = sentry_sdk.start_span( + op=op, + name=name, + origin=origin, + ) + for key, value in attributes.items(): + span.set_data(key, value) + + return span + except Exception: + pass + + return nullcontext() diff --git a/clients/python/src/taskbroker_client/worker/workerchild.py b/clients/python/src/taskbroker_client/worker/workerchild.py index 4f846ccb..6847b247 100644 --- a/clients/python/src/taskbroker_client/worker/workerchild.py +++ b/clients/python/src/taskbroker_client/worker/workerchild.py @@ -32,10 +32,12 @@ ) from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.crons import MonitorStatus, capture_checkin +from sentry_sdk.tracing import Span from taskbroker_client.app import import_app from taskbroker_client.constants import CompressionType from taskbroker_client.retry import NoRetriesRemainingError +from taskbroker_client.sdk import start_span, start_transaction from taskbroker_client.state import clear_current_task, current_task, set_current_task from taskbroker_client.task import Task from taskbroker_client.types import ContextHook, InflightTaskActivation, ProcessingResult @@ -612,45 +614,49 @@ def _execute_activation( kwargs = parameters.get("kwargs", {}) headers = dict(activation.headers) - transaction = sentry_sdk.continue_trace( - environ_or_headers=headers, - op="queue.task.taskworker", - name=activation.taskname, - origin="taskworker", - ) - sampling_context = { - "taskworker": { - "task": activation.taskname, - } - } + with ( metrics.track_memory_usage( "taskworker.worker.memory_change", tags={"namespace": activation.namespace, "taskname": activation.taskname}, ), sentry_sdk.isolation_scope(), - sentry_sdk.start_transaction(transaction, custom_sampling_context=sampling_context), + start_transaction( + name=activation.taskname, + op="queue.task.taskworker", + origin="taskworker", + attributes={ + "taskworker-task.id": activation.id, + }, + headers=headers, + sampling_context={ + "taskworker": { + "task": activation.taskname, + } + }, + ) as transaction, ): - transaction.set_data( - "taskworker-task", {"args": args, "kwargs": kwargs, "id": activation.id} - ) + # Do not attach on StreamedSpan because eager serialization increases memory use. + if isinstance(transaction, Span): + transaction.set_data("taskworker-task.args", args) + transaction.set_data("taskworker-task.kwargs", kwargs) + task_added_time = activation.received_at.ToDatetime().timestamp() # latency attribute needs to be in milliseconds latency = (time.time() - task_added_time) * 1000 - with sentry_sdk.start_span( - op=OP.QUEUE_PROCESS, + with start_span( name=activation.taskname, + op=OP.QUEUE_PROCESS, origin="taskworker", - ) as span: - span.set_data(SPANDATA.MESSAGING_DESTINATION_NAME, activation.namespace) - span.set_data(SPANDATA.MESSAGING_MESSAGE_ID, activation.id) - span.set_data(SPANDATA.MESSAGING_MESSAGE_RECEIVE_LATENCY, latency) - span.set_data( - SPANDATA.MESSAGING_MESSAGE_RETRY_COUNT, activation.retry_state.attempts - ) - span.set_data(SPANDATA.MESSAGING_SYSTEM, "taskworker") - + attributes={ + SPANDATA.MESSAGING_DESTINATION_NAME: activation.namespace, + SPANDATA.MESSAGING_MESSAGE_ID: activation.id, + SPANDATA.MESSAGING_MESSAGE_RECEIVE_LATENCY: latency, + SPANDATA.MESSAGING_MESSAGE_RETRY_COUNT: activation.retry_state.attempts, + SPANDATA.MESSAGING_SYSTEM: "taskworker", + }, + ): # TODO(taskworker) remove this when doing cleanup # The `__start_time` parameter is spliced into task parameters by # sentry.celery.SentryTask._add_metadata and needs to be removed diff --git a/uv.lock b/uv.lock index b1d03f28..674fbe70 100644 --- a/uv.lock +++ b/uv.lock @@ -703,14 +703,14 @@ wheels = [ [[package]] name = "sentry-sdk" -version = "2.49.0" +version = "2.66.1" source = { registry = "https://pypi.devinfra.sentry.io/simple" } dependencies = [ { name = "certifi", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "urllib3", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, ] wheels = [ - { url = "https://pypi.devinfra.sentry.io/wheels/sentry_sdk-2.49.0-py2.py3-none-any.whl", hash = "sha256:6ea78499133874445a20fe9c826c9e960070abeb7ae0cdf930314ab16bb97aa0" }, + { url = "https://pypi.devinfra.sentry.io/wheels/sentry_sdk-2.66.1-py3-none-any.whl", hash = "sha256:86002793161d9a95ef04bdd8d442e9bfece5d989b755f05d6360215094a7aff6" }, ] [package.optional-dependencies] @@ -856,7 +856,7 @@ requires-dist = [ { name = "redis-py-cluster", marker = "extra == 'cluster'", specifier = ">=2.1.0" }, { name = "sentry-arroyo", specifier = ">=2.41.0" }, { name = "sentry-protos", specifier = ">=0.26.1" }, - { name = "sentry-sdk", extras = ["http2"], specifier = ">=2.43.0" }, + { name = "sentry-sdk", extras = ["http2"], specifier = ">=2.66.1" }, { name = "setuptools", marker = "extra == 'examples'", specifier = ">=80.0" }, { name = "zstandard", specifier = ">=0.18.0" }, ] From d6885e4b35b7ef27d850397ee46ce48445274b81 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Mon, 27 Jul 2026 19:39:27 +0200 Subject: [PATCH 2/2] Add early return for when there is no active span --- clients/python/src/taskbroker_client/sdk.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/clients/python/src/taskbroker_client/sdk.py b/clients/python/src/taskbroker_client/sdk.py index 609360e2..e309f57c 100644 --- a/clients/python/src/taskbroker_client/sdk.py +++ b/clients/python/src/taskbroker_client/sdk.py @@ -57,6 +57,12 @@ def start_span( """Start a span in the currently active trace lifecycle.""" try: is_span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) + + # Early return to avoid large increase in span volume. + # Mirrors transaction-based tracing, in which `start_span()` no-ops when there is no active transaction. + if is_span_streaming and sentry_sdk.traces.get_current_span() is None: + return nullcontext() + if is_span_streaming: return sentry_sdk.traces.start_span( name=name,