Merge branch 'envelope-printer-transport-zcgvd' of github.com:getsentry/sentry-python into envelope-printer-transport-zcgvd

Author: ericapisani

.github/workflows/update-tox.yml

@@ -9,6 +9,7 @@ on:
 jobs:
   update-tox:
     name: Update test matrix
+    if: github.ref == 'refs/heads/master'
     runs-on: ubuntu-latest
     timeout-minutes: 10
 
@@ -59,7 +60,6 @@ jobs:
           BRANCH_NAME: ${{ steps.create-branch.outputs.branch_name }}
           COMMIT_TITLE: ${{ steps.create-branch.outputs.commit_title }}
           DATE: ${{ steps.create-branch.outputs.date }}
-          BASE_BRANCH: ${{ github.ref_name }}
         with:
           script: |
             const branchName = process.env.BRANCH_NAME;
@@ -105,7 +105,7 @@ jobs:
                 repo: context.repo.repo,
                 title: commitTitle + ' (' + date + ')',
                 head: branchName,
-                base: process.env.BASE_BRANCH,
+                base: 'master',
                 body: prBody,
             });
 

scripts/populate_tox/package_dependencies.jsonl

@@ -3,6 +3,7 @@
 import threading
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, TypeVar, Generic
+import weakref
 
 from sentry_sdk.utils import format_timestamp
 from sentry_sdk.envelope import Envelope, Item, PayloadRef
@@ -38,6 +39,26 @@ def __init__(
         self._flusher: "Optional[threading.Thread]" = None
         self._flusher_pid: "Optional[int]" = None
 
+        # See https://github.com/getsentry/sentry-python/blob/051cc01640a29bfd64b1f1e2e3414c02f027dd1b/sentry_sdk/monitor.py#L41-L50
+        if hasattr(os, "register_at_fork"):
+            weak_reset = weakref.WeakMethod(self._reset_thread_state)
+
+            def _reset_in_child() -> None:
+                method = weak_reset()
+                if method is not None:
+                    method()
+
+            os.register_at_fork(after_in_child=_reset_in_child)
+
+    def _reset_thread_state(self) -> None:
+        self._buffer = []
+        self._running = True
+        self._lock = threading.Lock()
+        self._active = threading.local()
+        self._flush_event = threading.Event()
+        self._flusher = None
+        self._flusher_pid = None
+
     def _ensure_thread(self) -> bool:
         """For forking processes we might need to restart this thread.
         This ensures that our process actually has that thread running.

scripts/populate_tox/releases.jsonl

@@ -1,7 +1,9 @@
-import threading
 from collections import defaultdict
 from datetime import datetime, timezone
+import os
+import threading
 from typing import TYPE_CHECKING
+import weakref
 
 from sentry_sdk._batcher import Batcher
 from sentry_sdk.envelope import Envelope, Item, PayloadRef
@@ -50,6 +52,30 @@ def __init__(
         self._flusher: "Optional[threading.Thread]" = None
         self._flusher_pid: "Optional[int]" = None
 
+        # See https://github.com/getsentry/sentry-python/blob/051cc01640a29bfd64b1f1e2e3414c02f027dd1b/sentry_sdk/monitor.py#L41-L50
+        if hasattr(os, "register_at_fork"):
+            weak_reset = weakref.WeakMethod(self._reset_thread_state)
+
+            def _reset_in_child() -> None:
+                method = weak_reset()
+                if method is not None:
+                    method()
+
+            os.register_at_fork(after_in_child=_reset_in_child)
+
+    def _reset_thread_state(self) -> None:
+        self._span_buffer = defaultdict(list)
+        self._running_size = defaultdict(lambda: 0)
+        self._running = True
+
+        self._lock = threading.Lock()
+        self._active = threading.local()
+
+        self._flush_event = threading.Event()
+
+        self._flusher = None
+        self._flusher_pid = None
+
     def add(self, span: "StreamedSpan") -> None:
         # Bail out if the current thread is already executing batcher code.
         # This prevents deadlocks when code running inside the batcher (e.g.

sentry_sdk/_batcher.py

@@ -6,7 +6,7 @@
 from sentry_sdk._init_implementation import init
 from sentry_sdk.consts import INSTRUMENTER
 from sentry_sdk.scope import Scope, _ScopeManager, new_scope, isolation_scope
-from sentry_sdk.traces import StreamedSpan
+from sentry_sdk.traces import StreamedSpan, _get_current_streamed_span
 from sentry_sdk.tracing import NoOpSpan, Transaction, trace
 from sentry_sdk.crons import monitor
 
@@ -422,7 +422,7 @@ def set_measurement(name: str, value: float, unit: "MeasurementUnit" = "") -> No
 
 def get_current_span(
     scope: "Optional[Scope]" = None,
-) -> "Optional[Union[Span, StreamedSpan]]":
+) -> "Optional[Span]":
     """
     Returns the currently active span if there is one running, otherwise `None`
     """
@@ -533,12 +533,7 @@ def update_current_span(
             attributes={"user_id": 123, "batch_size": 50}
         )
     """
-    current_span = get_current_span()
-
-    if current_span is None:
-        return
-
-    if isinstance(current_span, StreamedSpan):
+    if isinstance(_get_current_streamed_span(), StreamedSpan):
         warnings.warn(
             "The `update_current_span` API isn't available in streaming mode. "
             "Retrieve the current span with get_current_span() and use its API "
@@ -548,6 +543,11 @@ def update_current_span(
         )
         return
 
+    current_span = get_current_span()
+
+    if current_span is None:
+        return
+
     if op is not None:
         current_span.op = op
 

sentry_sdk/_span_batcher.py

@@ -15,7 +15,7 @@
 from sentry_sdk.integrations.celery.utils import _now_seconds_since_epoch
 from sentry_sdk.integrations.logging import ignore_logger
 from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.traces import StreamedSpan
+from sentry_sdk.traces import StreamedSpan, _get_current_streamed_span
 from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, Span, TransactionSource
 from sentry_sdk.tracing_utils import Baggage, has_span_streaming_enabled
 from sentry_sdk.utils import (
@@ -98,13 +98,16 @@ def setup_once() -> None:
 
 
 def _set_status(status: str) -> None:
+    client = sentry_sdk.get_client()
+    span_streaming = has_span_streaming_enabled(client.options)
+
     with capture_internal_exceptions():
         scope = sentry_sdk.get_current_scope()
-        if scope.span is not None:
-            if isinstance(scope.span, Span):
-                scope.span.set_status(status)
-            else:
-                scope.span.status = "ok" if status == "ok" else "error"
+
+        if span_streaming and scope.streamed_span is not None:
+            scope.streamed_span.status = "ok" if status == "ok" else "error"
+        elif not span_streaming and scope.span is not None:
+            scope.span.set_status(status)
 
 
 def _capture_exception(task: "Any", exc_info: "ExcInfo") -> None:
@@ -289,7 +292,7 @@ def apply_async(*args: "Any", **kwargs: "Any") -> "Any":
 
         span_mgr: "Union[StreamedSpan, Span, NoOpMgr]" = NoOpMgr()
         if span_streaming:
-            if not task_started_from_beat and sentry_sdk.get_current_span() is not None:
+            if not task_started_from_beat and _get_current_streamed_span() is not None:
                 span_mgr = sentry_sdk.traces.start_span(
                     name=task_name,
                     attributes={
@@ -570,7 +573,7 @@ def sentry_publish(self: "Producer", *args: "Any", **kwargs: "Any") -> "Any":
 
         span: "Union[StreamedSpan, Span, None]" = None
         if span_streaming:
-            if sentry_sdk.get_current_span() is not None:
+            if _get_current_streamed_span() is not None:
                 span = sentry_sdk.traces.start_span(
                     name=task_name,
                     attributes={

sentry_sdk/api.py

@@ -90,13 +90,17 @@ def _sentry_get_request_handler(*args: "Any", **kwargs: "Any") -> "Any":
             @wraps(old_call)
             def _sentry_call(*args: "Any", **kwargs: "Any") -> "Any":
                 current_scope = sentry_sdk.get_current_scope()
-                current_span = current_scope.span
 
-                if isinstance(current_span, StreamedSpan) and not isinstance(
-                    current_span, NoOpStreamedSpan
-                ):
-                    segment = current_span._segment
-                    segment._update_active_thread()
+                client = sentry_sdk.get_client()
+                if has_span_streaming_enabled(client.options):
+                    current_span = current_scope.streamed_span
+
+                    if isinstance(current_span, StreamedSpan) and not isinstance(
+                        current_span, NoOpStreamedSpan
+                    ):
+                        segment = current_span._segment
+                        segment._update_active_thread()
+
                 elif current_scope.transaction is not None:
                     current_scope.transaction.update_active_thread()
 

sentry_sdk/integrations/celery/__init__.py

@@ -21,7 +21,7 @@
 )
 from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
 from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.traces import NoOpStreamedSpan, StreamedSpan
+from sentry_sdk.traces import NoOpStreamedSpan, StreamedSpan, _get_current_streamed_span
 from sentry_sdk.tracing import (
     SOURCE_FOR_STYLE,
     TransactionSource,
@@ -254,7 +254,7 @@ def _default(value: "Any") -> "Any":
 def _set_request_body_data_on_streaming_segment(
     info: "Optional[Dict[str, Any]]",
 ) -> None:
-    current_span = sentry_sdk.get_current_span()
+    current_span = _get_current_streamed_span()
     if (
         info
         and "data" in info
@@ -545,19 +545,22 @@ def event_processor(
 
             @functools.wraps(old_func)
             def _sentry_sync_func(*args: "Any", **kwargs: "Any") -> "Any":
-                integration = sentry_sdk.get_client().get_integration(
-                    StarletteIntegration
-                )
+                client = sentry_sdk.get_client()
+
+                integration = client.get_integration(StarletteIntegration)
                 if integration is None:
                     return old_func(*args, **kwargs)
 
                 current_scope = sentry_sdk.get_current_scope()
-                current_span = current_scope.span
 
-                if isinstance(current_span, StreamedSpan) and not isinstance(
-                    current_span, NoOpStreamedSpan
-                ):
-                    current_span._segment._update_active_thread()
+                span_streaming = has_span_streaming_enabled(client.options)
+                if span_streaming:
+                    current_span = current_scope.streamed_span
+
+                    if isinstance(current_span, StreamedSpan) and not isinstance(
+                        current_span, NoOpStreamedSpan
+                    ):
+                        current_span._segment._update_active_thread()
                 elif current_scope.transaction is not None:
                     current_scope.transaction.update_active_thread()
 

sentry_sdk/integrations/fastapi.py

@@ -580,12 +580,18 @@ def get_traceparent(self, *args: "Any", **kwargs: "Any") -> "Optional[str]":
         """
         client = self.get_client()
 
+        if not has_tracing_enabled(client.options):
+            return self.get_active_propagation_context().to_traceparent()
+
+        span_streaming = has_span_streaming_enabled(client.options)
         # If we have an active span, return traceparent from there
         if (
-            has_tracing_enabled(client.options)
-            and self.span is not None
-            and not isinstance(self.span, NoOpStreamedSpan)
+            span_streaming
+            and self.streamed_span is not None
+            and not isinstance(self.streamed_span, NoOpStreamedSpan)
         ):
+            return self.streamed_span._to_traceparent()
+        elif not span_streaming and self.span is not None:
             return self.span._to_traceparent()
 
         # else return traceparent from the propagation context
@@ -598,12 +604,18 @@ def get_baggage(self, *args: "Any", **kwargs: "Any") -> "Optional[Baggage]":
         """
         client = self.get_client()
 
+        if not has_tracing_enabled(client.options):
+            return self.get_active_propagation_context().get_baggage()
+
+        span_streaming = has_span_streaming_enabled(client.options)
         # If we have an active span, return baggage from there
         if (
-            has_tracing_enabled(client.options)
-            and self.span is not None
-            and not isinstance(self.span, NoOpStreamedSpan)
+            span_streaming
+            and self.streamed_span is not None
+            and not isinstance(self.streamed_span, NoOpStreamedSpan)
         ):
+            return self.streamed_span._to_baggage()
+        elif not span_streaming and self.span is not None:
             return self.span._to_baggage()
 
         # else return baggage from the propagation context
@@ -680,7 +692,9 @@ def iter_trace_propagation_headers(
             return
 
         span = kwargs.pop("span", None)
-        span = span or self.span
+        if not span:
+            span_streaming = has_span_streaming_enabled(client.options)
+            span = self.streamed_span if span_streaming else self.span
 
         if (
             has_tracing_enabled(client.options)
@@ -877,12 +891,12 @@ def set_user(self, value: "Optional[Dict[str, Any]]") -> None:
             session.update(user=value)
 
     @property
-    def span(self) -> "Optional[Union[Span, StreamedSpan]]":
+    def span(self) -> "Optional[Span]":
         """Get/set current tracing span or transaction."""
-        return self._span
+        return self._span if isinstance(self._span, Span) else None
 
     @span.setter
-    def span(self, span: "Optional[Union[Span, StreamedSpan]]") -> None:
+    def span(self, span: "Optional[Span]") -> None:
         self._span = span
         # XXX: this differs from the implementation in JS, there Scope.setSpan
         # does not set Scope._transactionName.
@@ -893,6 +907,15 @@ def span(self, span: "Optional[Union[Span, StreamedSpan]]") -> None:
                 if transaction.source:
                     self._transaction_info["source"] = transaction.source
 
+    @property
+    def streamed_span(self) -> "Optional[StreamedSpan]":
+        """Get/set current tracing span."""
+        return self._span if isinstance(self._span, StreamedSpan) else None
+
+    @streamed_span.setter
+    def streamed_span(self, span: "Optional[StreamedSpan]") -> None:
+        self._span = span
+
         # Also set _transaction and _transaction_info in streaming mode as this
         # is used for populating events and linking them to segments
         if (
@@ -1267,7 +1290,7 @@ def start_streamed_span(
         if parent_span is _DEFAULT_PARENT_SPAN or isinstance(
             parent_span, NoOpStreamedSpan
         ):
-            parent_span = self.span  # type: ignore
+            parent_span = self.streamed_span
 
         # If no eligible parent_span was provided and there is no currently
         # active span, this is a segment
@@ -1848,9 +1871,21 @@ def apply_to_telemetry(self, telemetry: "Union[Log, Metric, StreamedSpan]") -> N
                 telemetry["trace_id"] = (
                     trace_id or "00000000-0000-0000-0000-000000000000"
                 )
-            span_id = trace_context.get("span_id")
-            if telemetry.get("span_id") is None and span_id:
-                telemetry["span_id"] = span_id
+
+            # span_id should only be populated if there's an active span. We can't
+            # use the trace_context here because it synthesizes a span_id if there
+            # isn't one
+            if telemetry.get("span_id") is None:
+                if self._span is not None and not isinstance(
+                    self._span, NoOpStreamedSpan
+                ):
+                    telemetry["span_id"] = self._span.span_id
+                else:
+                    external_propagation_context = get_external_propagation_context()
+                    if external_propagation_context:
+                        _, span_id = external_propagation_context
+                        if span_id is not None:
+                            telemetry["span_id"] = span_id
 
         self._apply_scope_attributes_to_telemetry(telemetry)
         self._apply_user_attributes_to_telemetry(telemetry)