Trace e2e

2 files changed, 44 insertions, 0 deletions

diff --git a/synapse/storage/e2e_room_keys.py b/synapse/storage/e2e_room_keys.py
index 99128f2df7..b0e426d30e 100644
--- a/synapse/storage/e2e_room_keys.py
+++ b/synapse/storage/e2e_room_keys.py
@@ -18,11 +18,13 @@ import json
 from twisted.internet import defer
 
 from synapse.api.errors import StoreError
+from synapse.util.tracerutils import trace_defered_function, trace_function
 
 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
@@ -63,6 +65,7 @@ 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
@@ -95,6 +98,7 @@ class EndToEndRoomKeyStore(SQLBaseStore):
             lock=False,
         )
 
+    @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
@@ -153,6 +157,7 @@ class EndToEndRoomKeyStore(SQLBaseStore):
 
         return sessions
 
+    @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
@@ -194,6 +199,7 @@ class EndToEndRoomKeyStore(SQLBaseStore):
             raise StoreError(404, "No current backup version")
         return row[0]
 
+    @trace_function
     def get_e2e_room_keys_version_info(self, user_id, version=None):
         """Get info metadata about a version of our room_keys backup.
 
@@ -236,6 +242,7 @@ class EndToEndRoomKeyStore(SQLBaseStore):
             "get_e2e_room_keys_version_info", _get_e2e_room_keys_version_info_txn
         )
 
+    @trace_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.
@@ -276,6 +283,7 @@ class EndToEndRoomKeyStore(SQLBaseStore):
             "create_e2e_room_keys_version_txn", _create_e2e_room_keys_version_txn
         )
 
+    @trace_function
     def update_e2e_room_keys_version(self, user_id, version, info):
         """Update a given backup version
 
@@ -292,6 +300,7 @@ class EndToEndRoomKeyStore(SQLBaseStore):
             desc="update_e2e_room_keys_version",
         )
 
+    @trace_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.
diff --git a/synapse/storage/end_to_end_keys.py b/synapse/storage/end_to_end_keys.py
index 1e07474e70..9779cb70c9 100644
--- a/synapse/storage/end_to_end_keys.py
+++ b/synapse/storage/end_to_end_keys.py
@@ -22,8 +22,11 @@ from synapse.util.caches.descriptors import cached
 
 from ._base import SQLBaseStore, db_to_json
 
+from synapse.util.tracerutils import TracerUtil, trace_defered_function, trace_function
+
 
 class EndToEndKeyWorkerStore(SQLBaseStore):
+    @trace_defered_function
     @defer.inlineCallbacks
     def get_e2e_device_keys(
         self, query_list, include_all_devices=False, include_deleted_devices=False
@@ -40,6 +43,7 @@ class EndToEndKeyWorkerStore(SQLBaseStore):
             Dict mapping from user-id to dict mapping from device_id to
             dict containing "key_json", "device_display_name".
         """
+        TracerUtil.set_tag("query_list", query_list)
         if not query_list:
             return {}
 
@@ -57,9 +61,13 @@ class EndToEndKeyWorkerStore(SQLBaseStore):
 
         return results
 
+    @trace_function
     def _get_e2e_device_keys_txn(
         self, txn, query_list, include_all_devices=False, include_deleted_devices=False
     ):
+        TracerUtil.set_tag("include_all_devices", include_all_devices)
+        TracerUtil.set_tag("include_deleted_devices", include_deleted_devices)
+
         query_clauses = []
         query_params = []
 
@@ -104,8 +112,10 @@ class EndToEndKeyWorkerStore(SQLBaseStore):
             for user_id, device_id in deleted_devices:
                 result.setdefault(user_id, {})[device_id] = None
 
+        TracerUtil.log_kv(result)
         return result
 
+    @trace_defered_function
     @defer.inlineCallbacks
     def get_e2e_one_time_keys(self, user_id, device_id, key_ids):
         """Retrieve a number of one-time keys for a user
@@ -121,6 +131,10 @@ class EndToEndKeyWorkerStore(SQLBaseStore):
             key_id) to json string for key
         """
 
+        TracerUtil.set_tag("user_id", user_id)
+        TracerUtil.set_tag("device_id", device_id)
+        TracerUtil.set_tag("key_ids", key_ids)
+
         rows = yield self._simple_select_many_batch(
             table="e2e_one_time_keys_json",
             column="key_id",
@@ -145,7 +159,11 @@ class EndToEndKeyWorkerStore(SQLBaseStore):
                 (algorithm, key_id, key json)
         """
 
+        @trace_function
         def _add_e2e_one_time_keys(txn):
+            TracerUtil.set_tag("user_id", user_id)
+            TracerUtil.set_tag("device_id", device_id)
+            TracerUtil.set_tag("new_keys", new_keys)
             # We are protected from race between lookup and insertion due to
             # a unique constraint. If there is a race of two calls to
             # `add_e2e_one_time_keys` then they'll conflict and we will only
@@ -201,7 +219,13 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
         or the keys were already in the database.
         """
 
+        @trace_function
         def _set_e2e_device_keys_txn(txn):
+            TracerUtil.set_tag("user_id", user_id)
+            TracerUtil.set_tag("device_id", device_id)
+            TracerUtil.set_tag("time_now", time_now)
+            TracerUtil.set_tag("device_keys", device_keys)
+
             old_key_json = self._simple_select_one_onecol_txn(
                 txn,
                 table="e2e_device_keys_json",
@@ -215,6 +239,7 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
             new_key_json = encode_canonical_json(device_keys).decode("utf-8")
 
             if old_key_json == new_key_json:
+                TracerUtil.set_tag("error", True)
                 return False
 
             self._simple_upsert_txn(
@@ -231,6 +256,7 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
     def claim_e2e_one_time_keys(self, query_list):
         """Take a list of one time keys out of the database"""
 
+        @trace_function
         def _claim_e2e_one_time_keys(txn):
             sql = (
                 "SELECT key_id, key_json FROM e2e_one_time_keys_json"
@@ -252,7 +278,13 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
                 " AND key_id = ?"
             )
             for user_id, device_id, algorithm, key_id in delete:
+                TracerUtil.log_kv(
+                    {"message": "executing claim transaction on database"}
+                )
                 txn.execute(sql, (user_id, device_id, algorithm, key_id))
+                TracerUtil.log_kv(
+                    {"message": "finished executing and invalidating cache"}
+                )
                 self._invalidate_cache_and_stream(
                     txn, self.count_e2e_one_time_keys, (user_id, device_id)
                 )
@@ -261,7 +293,10 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
         return self.runInteraction("claim_e2e_one_time_keys", _claim_e2e_one_time_keys)
 
     def delete_e2e_keys_by_device(self, user_id, device_id):
+        @trace_function
         def delete_e2e_keys_by_device_txn(txn):
+            TracerUtil.set_tag("user_id", user_id)
+            TracerUtil.set_tag("device_id", device_id)
             self._simple_delete_txn(
                 txn,
                 table="e2e_device_keys_json",