Add opentracing spans to DB calls and measure blocks github/erikj/opentracing_db_measure erikj/opentracing_db_measure

2 files changed, 23 insertions, 12 deletions

diff --git a/synapse/storage/database.py b/synapse/storage/database.py
index bd39c095af..f496ab35ed 100644
--- a/synapse/storage/database.py
+++ b/synapse/storage/database.py
@@ -45,6 +45,7 @@ from synapse.logging.context import (
     current_context,
     make_deferred_yieldable,
 )
+from synapse.logging.opentracing import start_active_span
 from synapse.metrics.background_process_metrics import run_as_background_process
 from synapse.storage.background_updates import BackgroundUpdater
 from synapse.storage.engines import BaseDatabaseEngine, PostgresEngine, Sqlite3Engine
@@ -653,19 +654,20 @@ class DatabasePool:
             logger.warning("Starting db txn '%s' from sentinel context", desc)
 
         try:
-            result = await self.runWithConnection(
-                self.new_transaction,
-                desc,
-                after_callbacks,
-                exception_callbacks,
-                func,
-                *args,
-                db_autocommit=db_autocommit,
-                **kwargs,
-            )
+            with start_active_span(f"txn.{desc}"):
+                result = await self.runWithConnection(
+                    self.new_transaction,
+                    desc,
+                    after_callbacks,
+                    exception_callbacks,
+                    func,
+                    *args,
+                    db_autocommit=db_autocommit,
+                    **kwargs,
+                )
 
-            for after_callback, after_args, after_kwargs in after_callbacks:
-                after_callback(*after_args, **after_kwargs)
+                for after_callback, after_args, after_kwargs in after_callbacks:
+                    after_callback(*after_args, **after_kwargs)
         except Exception:
             for after_callback, after_args, after_kwargs in exception_callbacks:
                 after_callback(*after_args, **after_kwargs)
diff --git a/synapse/util/metrics.py b/synapse/util/metrics.py
index 6d14351bd2..c6ff520b34 100644
--- a/synapse/util/metrics.py
+++ b/synapse/util/metrics.py
@@ -23,6 +23,7 @@ from synapse.logging.context import (
     LoggingContext,
     current_context,
 )
+from synapse.logging.opentracing import start_active_span
 from synapse.metrics import InFlightGauge
 
 logger = logging.getLogger(__name__)
@@ -102,6 +103,7 @@ class Measure:
         "name",
         "_logging_context",
         "start",
+        "_span",
     ]
 
     def __init__(self, clock, name: str):
@@ -126,13 +128,18 @@ class Measure:
         self._logging_context = LoggingContext(str(curr_context), parent_context)
         self.start = None  # type: Optional[int]
 
+        self._span = start_active_span(f"measure.{name}")
+
     def __enter__(self) -> "Measure":
         if self.start is not None:
             raise RuntimeError("Measure() objects cannot be re-used")
 
+        self._span.__enter__()
+
         self.start = self.clock.time()
         self._logging_context.__enter__()
         in_flight.register((self.name,), self._update_in_flight)
+
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -156,6 +163,8 @@ class Measure:
         except ValueError:
             logger.warning("Failed to save metrics! Usage: %s", usage)
 
+        self._span.__exit__(exc_type, exc_val, exc_tb)
+
     def get_resource_usage(self) -> ContextResourceUsage:
         """Get the resources used within this Measure block