From 6255a1a6222b93f21a30fcf4d0d1c4e63c8b44f5 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <erice@element.io>
Date: Mon, 1 Aug 2022 19:05:03 -0500
Subject: [PATCH] Fix tests and some lints

---
 .../federation/sender/transaction_manager.py  | 13 ++++++++----
 synapse/federation/units.py                   |  4 ++--
 synapse/logging/tracing.py                    | 20 ++++++++++++++++++-
 tests/logging/test_tracing.py                 |  2 +-
 4 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/synapse/federation/sender/transaction_manager.py b/synapse/federation/sender/transaction_manager.py
index 38559a2baf..c819cdcf8b 100644
--- a/synapse/federation/sender/transaction_manager.py
+++ b/synapse/federation/sender/transaction_manager.py
@@ -25,6 +25,7 @@ from synapse.logging.tracing import (
     Link,
     StatusCode,
     extract_text_map,
+    get_span_context_from_context,
     set_status,
     start_active_span,
     whitelisted_homeserver,
@@ -89,11 +90,15 @@ class TransactionManager:
         keep_destination = whitelisted_homeserver(destination)
 
         for edu in edus:
-            context = edu.get_context()
-            if context:
-                span_contexts.append(extract_text_map(json_decoder.decode(context)))
+            tracing_context_json = edu.get_tracing_context_json()
+            if tracing_context_json:
+                context = extract_text_map(json_decoder.decode(tracing_context_json))
+                if context:
+                    span_context = get_span_context_from_context(context)
+                    if span_context:
+                        span_contexts.append(span_context)
             if keep_destination:
-                edu.strip_context()
+                edu.strip_tracing_context()
 
         with start_active_span(
             "send_transaction",
diff --git a/synapse/federation/units.py b/synapse/federation/units.py
index 9d6295f552..a6b590269e 100644
--- a/synapse/federation/units.py
+++ b/synapse/federation/units.py
@@ -55,12 +55,12 @@ class Edu:
             "destination": self.destination,
         }
 
-    def get_context(self) -> str:
+    def get_tracing_context_json(self) -> str:
         return getattr(self, "content", {}).get(
             EventContentFields.TRACING_CONTEXT, "{}"
         )
 
-    def strip_context(self) -> None:
+    def strip_tracing_context(self) -> None:
         getattr(self, "content", {})[EventContentFields.TRACING_CONTEXT] = "{}"
 
 
diff --git a/synapse/logging/tracing.py b/synapse/logging/tracing.py
index be0bfe2d18..d76db0fbc1 100644
--- a/synapse/logging/tracing.py
+++ b/synapse/logging/tracing.py
@@ -310,7 +310,7 @@ R = TypeVar("R")
 
 
 def only_if_tracing(func: Callable[P, R]) -> Callable[P, Optional[R]]:
-    """Executes the function only if we're tracing. Otherwise returns None."""
+    """Decorator function that executes the function only if we're tracing. Otherwise returns None."""
 
     @wraps(func)
     def _only_if_tracing_inner(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
@@ -544,8 +544,12 @@ def start_active_span(
         start_time=start_time,
         record_exception=record_exception,
         set_status_on_exception=set_status_on_exception,
+        tracer=tracer,
     )
 
+    ctx = opentelemetry.trace.propagation.set_span_in_context(span)
+    logger.info("efwfewaafwewffew ctx=%s span_context=%s", ctx, span.get_span_context())
+
     # Equivalent to `tracer.start_as_current_span`
     return opentelemetry.trace.use_span(
         span,
@@ -587,6 +591,20 @@ def get_active_span() -> Optional["opentelemetry.trace.span.Span"]:
     return opentelemetry.trace.get_current_span()
 
 
+def get_span_context_from_context(
+    context: "opentelemetry.context.context.Context",
+) -> Optional["opentelemetry.trace.span.SpanContext"]:
+    """Utility function to convert a `Context` to a `SpanContext`
+
+    Based on https://github.com/open-telemetry/opentelemetry-python/blob/43288ca9a36144668797c11ca2654836ec8b5e99/opentelemetry-api/src/opentelemetry/trace/propagation/tracecontext.py#L99-L102
+    """
+    span = opentelemetry.trace.get_current_span(context=context)
+    span_context = span.get_span_context()
+    if span_context == opentelemetry.trace.INVALID_SPAN_CONTEXT:
+        return None
+    return span_context
+
+
 @ensure_active_span("set a tag")
 def set_attribute(key: str, value: Union[str, bool, int, float]) -> None:
     """Sets a tag on the active span"""
diff --git a/tests/logging/test_tracing.py b/tests/logging/test_tracing.py
index 27b9d4348b..59529c9d01 100644
--- a/tests/logging/test_tracing.py
+++ b/tests/logging/test_tracing.py
@@ -32,7 +32,7 @@ except ImportError:
     opentelemetry = None  # type: ignore[assignment]
 
 
-class LogContextScopeManagerTestCase(TestCase):
+class TracingTestCase(TestCase):
     """
     Test logging contexts and active opentelemetry spans.
     """