From e07a8caf58fad2c56518560cfc31d90a761bd5a9 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Wed, 17 Jun 2020 14:13:41 +0100
Subject: [PATCH] Add support for using rust-python-jaeger-reporter (#7697)

---
 changelog.d/7697.misc          |  1 +
 mypy.ini                       |  3 +++
 synapse/logging/opentracing.py | 39 +++++++++++++++++++++++++++++++---
 3 files changed, 40 insertions(+), 3 deletions(-)
 create mode 100644 changelog.d/7697.misc

diff --git a/changelog.d/7697.misc b/changelog.d/7697.misc
new file mode 100644
index 000000000..345862b5a
--- /dev/null
+++ b/changelog.d/7697.misc
@@ -0,0 +1 @@
+Add support for using `rust-python-jaeger-reporter` library to reduce jaeger tracing overhead.
diff --git a/mypy.ini b/mypy.ini
index 3533797d6..a61009b19 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -78,3 +78,6 @@ ignore_missing_imports = True
 
 [mypy-authlib.*]
 ignore_missing_imports = True
+
+[mypy-rust_python_jaeger_reporter.*]
+ignore_missing_imports = True
diff --git a/synapse/logging/opentracing.py b/synapse/logging/opentracing.py
index 5dddf5700..73bef5e5c 100644
--- a/synapse/logging/opentracing.py
+++ b/synapse/logging/opentracing.py
@@ -171,8 +171,9 @@ import logging
 import re
 import types
 from functools import wraps
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING, Dict, Optional, Type
 
+import attr
 from canonicaljson import json
 
 from twisted.internet import defer
@@ -232,6 +233,30 @@ except ImportError:
     LogContextScopeManager = None  # type: ignore
 
 
+try:
+    from rust_python_jaeger_reporter import Reporter
+
+    @attr.s(slots=True, frozen=True)
+    class _WrappedRustReporter:
+        """Wrap the reporter to ensure `report_span` never throws.
+        """
+
+        _reporter = attr.ib(type=Reporter, default=attr.Factory(Reporter))
+
+        def set_process(self, *args, **kwargs):
+            return self._reporter.set_process(*args, **kwargs)
+
+        def report_span(self, span):
+            try:
+                return self._reporter.report_span(span)
+            except Exception:
+                logger.exception("Failed to report span")
+
+    RustReporter = _WrappedRustReporter  # type: Optional[Type[_WrappedRustReporter]]
+except ImportError:
+    RustReporter = None
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -320,11 +345,19 @@ def init_tracer(hs: "HomeServer"):
 
     set_homeserver_whitelist(hs.config.opentracer_whitelist)
 
-    JaegerConfig(
+    config = JaegerConfig(
         config=hs.config.jaeger_config,
         service_name="{} {}".format(hs.config.server_name, hs.get_instance_name()),
         scope_manager=LogContextScopeManager(hs.config),
-    ).initialize_tracer()
+    )
+
+    # If we have the rust jaeger reporter available let's use that.
+    if RustReporter:
+        logger.info("Using rust_python_jaeger_reporter library")
+        tracer = config.create_tracer(RustReporter(), config.sampler)
+        opentracing.set_global_tracer(tracer)
+    else:
+        config.initialize_tracer()
 
 
 # Whitelisting