ref(o11y): Support SDK configured with the streaming trace lifecycle

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",

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=[]),

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()

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