Link the send loop with the edus contexts

The contexts were being filtered too early so  the send loop wasn't
being linked to them unless the destination
was whitelisted.

3 files changed, 12 insertions, 7 deletions

diff --git a/synapse/federation/sender/transaction_manager.py b/synapse/federation/sender/transaction_manager.py
index 62ca6a3e87..42f46394bc 100644
--- a/synapse/federation/sender/transaction_manager.py
+++ b/synapse/federation/sender/transaction_manager.py
@@ -26,6 +26,7 @@ from synapse.logging.opentracing import (
     set_tag,
     start_active_span_follows_from,
     tags,
+    whitelisted_homeserver,
 )
 from synapse.util.metrics import measure_func
 
@@ -59,9 +60,13 @@ class TransactionManager(object):
         # The span_contexts is a generator so that it won't be evaluated if
         # opentracing is disabled. (Yay speed!)
 
-        span_contexts = (
-            extract_text_map(json.loads(edu.get_context())) for edu in pending_edus
-        )
+        span_contexts = []
+        keep_destination = whitelisted_homeserver(destination)
+
+        for edu in pending_edus:
+            span_contexts.append(extract_text_map(json.loads(edu.get_context())))
+            if keep_destination:
+                edu.strip_context()
 
         with start_active_span_follows_from("send_transaction", span_contexts):
 
diff --git a/synapse/federation/units.py b/synapse/federation/units.py
index aa84621206..b4d743cde7 100644
--- a/synapse/federation/units.py
+++ b/synapse/federation/units.py
@@ -41,6 +41,9 @@ class Edu(JsonEncodedObject):
     def get_context(self):
         return getattr(self, "content", {}).get("org.matrix.opentracing_context", "{}")
 
+    def strip_context(self):
+        getattr(self, "content", {})["org.matrix.opentracing_context"] = "{}"
+
 
 class Transaction(JsonEncodedObject):
     """ A transaction is a list of Pdus and Edus to be sent to a remote home
diff --git a/synapse/handlers/devicemessage.py b/synapse/handlers/devicemessage.py
index 01731cb2d0..0043cbea17 100644
--- a/synapse/handlers/devicemessage.py
+++ b/synapse/handlers/devicemessage.py
@@ -25,7 +25,6 @@ from synapse.logging.opentracing import (
     log_kv,
     set_tag,
     start_active_span,
-    whitelisted_homeserver,
 )
 from synapse.types import UserID, get_domain_from_id
 from synapse.util.stringutils import random_string
@@ -121,9 +120,7 @@ class DeviceMessageHandler(object):
                     "sender": sender_user_id,
                     "type": message_type,
                     "message_id": message_id,
-                    "org.matrix.opentracing_context": json.dumps(context)
-                    if whitelisted_homeserver(destination)
-                    else None,
+                    "org.matrix.opentracing_context": json.dumps(context),
                 }
 
         log_kv({"local_messages": local_messages})