Fix exemplars

3 files changed, 27 insertions, 7 deletions

diff --git a/synapse/http/request_metrics.py b/synapse/http/request_metrics.py
index 4acb30e1db..e7c672199a 100644
--- a/synapse/http/request_metrics.py
+++ b/synapse/http/request_metrics.py
@@ -16,7 +16,7 @@
 import logging
 import threading
 import traceback
-from typing import Dict, Mapping, Set, Tuple
+from typing import TYPE_CHECKING, Dict, Mapping, Optional, Set, Tuple
 
 from prometheus_client.core import Counter, Histogram
 
@@ -24,6 +24,9 @@ from synapse.logging.context import current_context
 from synapse.logging.opentracing import get_prometheus_exemplar
 from synapse.metrics import LaterGauge
 
+if TYPE_CHECKING:
+    import opentracing
+
 logger = logging.getLogger(__name__)
 
 
@@ -144,7 +147,12 @@ LaterGauge(
 
 
 class RequestMetrics:
-    def start(self, time_sec: float, name: str, method: str) -> None:
+    def start(
+        self,
+        time_sec: float,
+        name: str,
+        method: str,
+    ) -> None:
         self.start_ts = time_sec
         self.start_context = current_context()
         self.name = name
@@ -163,7 +171,13 @@ class RequestMetrics:
         with _in_flight_requests_lock:
             _in_flight_requests.add(self)
 
-    def stop(self, time_sec: float, response_code: int, sent_bytes: int) -> None:
+    def stop(
+        self,
+        time_sec: float,
+        response_code: int,
+        sent_bytes: int,
+        span: Optional["opentracing.Span"],
+    ) -> None:
         with _in_flight_requests_lock:
             _in_flight_requests.discard(self)
 
@@ -194,7 +208,7 @@ class RequestMetrics:
         response_count.labels(self.method, self.name, tag).inc()
 
         response_timer.labels(self.method, self.name, tag, response_code_str).observe(
-            time_sec - self.start_ts, exemplar=get_prometheus_exemplar()
+            time_sec - self.start_ts, exemplar=get_prometheus_exemplar(span)
         )
 
         resource_usage = context.get_resource_usage()
diff --git a/synapse/http/site.py b/synapse/http/site.py
index 5b5a7c1e59..82f6281141 100644
--- a/synapse/http/site.py
+++ b/synapse/http/site.py
@@ -487,7 +487,9 @@ class SynapseRequest(Request):
             self._opentracing_span.finish()
 
         try:
-            self.request_metrics.stop(self.finish_time, self.code, self.sentLength)
+            self.request_metrics.stop(
+                self.finish_time, self.code, self.sentLength, self._opentracing_span
+            )
         except Exception as e:
             logger.warning("Failed to stop metrics: %r", e)
 
diff --git a/synapse/logging/opentracing.py b/synapse/logging/opentracing.py
index f5e4993fab..9a7fd7026b 100644
--- a/synapse/logging/opentracing.py
+++ b/synapse/logging/opentracing.py
@@ -672,8 +672,12 @@ def active_span() -> Optional["opentracing.Span"]:
     return opentracing.tracer.active_span
 
 
-def get_prometheus_exemplar() -> Optional[Dict[str, str]]:
-    span = active_span()
+def get_prometheus_exemplar(
+    span: Optional["opentracing.Span"] = None,
+) -> Optional[Dict[str, str]]:
+    if not span:
+        span = active_span()
+
     if not span:
         return None