Clean up room key tracing

2 files changed, 32 insertions, 12 deletions

diff --git a/synapse/handlers/e2e_room_keys.py b/synapse/handlers/e2e_room_keys.py
index 0459d52a3c..1c34ff712f 100644
--- a/synapse/handlers/e2e_room_keys.py
+++ b/synapse/handlers/e2e_room_keys.py
@@ -179,7 +179,6 @@ class E2eRoomKeysHandler(object):
                         user_id, version, room_id, session_id, session
                     )
 
-    @opentracing.trace_defered_function
     @defer.inlineCallbacks
     def _upload_room_key(self, user_id, version, room_id, session_id, room_key):
         """Upload a given room_key for a given room and session into a given
@@ -192,7 +191,14 @@ class E2eRoomKeysHandler(object):
             session_id(str): the session whose room_key we're setting
             room_key(dict): the room_key being set
         """
-
+        opentracing.log_kv(
+            {
+                "message": "Trying to upload room key",
+                "room_id": room_id,
+                "session_id": session_id,
+                "user_id": user_id,
+            }
+        )
         # get the room_key for this particular row
         current_room_key = None
         try:
@@ -201,14 +207,23 @@ class E2eRoomKeysHandler(object):
             )
         except StoreError as e:
             if e.code == 404:
-                pass
+                opentracing.log_kv(
+                    {
+                        "message": "Room key not found.",
+                        "room_id": room_id,
+                        "user_id": user_id,
+                    }
+                )
             else:
                 raise
 
         if self._should_replace_room_key(current_room_key, room_key):
+            opentracing.log_kv({"message": "Replacing room key."})
             yield self.store.set_e2e_room_key(
                 user_id, version, room_id, session_id, room_key
             )
+        else:
+            opentracing.log_kv({"message": "Not replacing room_key."})
 
     @staticmethod
     def _should_replace_room_key(current_room_key, room_key):
diff --git a/synapse/storage/e2e_room_keys.py b/synapse/storage/e2e_room_keys.py
index f18605dfc4..bfd8b51362 100644
--- a/synapse/storage/e2e_room_keys.py
+++ b/synapse/storage/e2e_room_keys.py
@@ -18,13 +18,12 @@ import json
 from twisted.internet import defer
 
 from synapse.api.errors import StoreError
-from synapse.logging.opentracing import trace_defered_function
+import synapse.logging.opentracing as opentracing
 
 from ._base import SQLBaseStore
 
 
 class EndToEndRoomKeyStore(SQLBaseStore):
-    @trace_defered_function
     @defer.inlineCallbacks
     def get_e2e_room_key(self, user_id, version, room_id, session_id):
         """Get the encrypted E2E room key for a given session from a given
@@ -65,7 +64,6 @@ class EndToEndRoomKeyStore(SQLBaseStore):
 
         return row
 
-    @trace_defered_function
     @defer.inlineCallbacks
     def set_e2e_room_key(self, user_id, version, room_id, session_id, room_key):
         """Replaces or inserts the encrypted E2E room key for a given session in
@@ -97,8 +95,16 @@ class EndToEndRoomKeyStore(SQLBaseStore):
             },
             lock=False,
         )
+        opentracing.log_kv(
+            {
+                "message": "Set room key",
+                "room_id": room_id,
+                "session_id": session_id,
+                "room_key": room_key,
+            }
+        )
 
-    @trace_defered_function
+    @opentracing.trace_defered_function
     @defer.inlineCallbacks
     def get_e2e_room_keys(self, user_id, version, room_id=None, session_id=None):
         """Bulk get the E2E room keys for a given backup, optionally filtered to a given
@@ -157,7 +163,7 @@ class EndToEndRoomKeyStore(SQLBaseStore):
 
         return sessions
 
-    @trace_defered_function
+    @opentracing.trace_defered_function
     @defer.inlineCallbacks
     def delete_e2e_room_keys(self, user_id, version, room_id=None, session_id=None):
         """Bulk delete the E2E room keys for a given backup, optionally filtered to a given
@@ -199,7 +205,6 @@ class EndToEndRoomKeyStore(SQLBaseStore):
             raise StoreError(404, "No current backup version")
         return row[0]
 
-    @trace_defered_function
     def get_e2e_room_keys_version_info(self, user_id, version=None):
         """Get info metadata about a version of our room_keys backup.
 
@@ -242,7 +247,7 @@ class EndToEndRoomKeyStore(SQLBaseStore):
             "get_e2e_room_keys_version_info", _get_e2e_room_keys_version_info_txn
         )
 
-    @trace_defered_function
+    @opentracing.trace_defered_function
     def create_e2e_room_keys_version(self, user_id, info):
         """Atomically creates a new version of this user's e2e_room_keys store
         with the given version info.
@@ -283,7 +288,7 @@ class EndToEndRoomKeyStore(SQLBaseStore):
             "create_e2e_room_keys_version_txn", _create_e2e_room_keys_version_txn
         )
 
-    @trace_defered_function
+    @opentracing.trace_defered_function
     def update_e2e_room_keys_version(self, user_id, version, info):
         """Update a given backup version
 
@@ -300,7 +305,7 @@ class EndToEndRoomKeyStore(SQLBaseStore):
             desc="update_e2e_room_keys_version",
         )
 
-    @trace_defered_function
+    @opentracing.trace_defered_function
     def delete_e2e_room_keys_version(self, user_id, version=None):
         """Delete a given backup version of the user's room keys.
         Doesn't delete their actual key data.