Cleanup key upload tracing

2 files changed, 30 insertions, 16 deletions

diff --git a/synapse/handlers/e2e_keys.py b/synapse/handlers/e2e_keys.py
index 01d5d383bf..a4b2f18e6b 100644
--- a/synapse/handlers/e2e_keys.py
+++ b/synapse/handlers/e2e_keys.py
@@ -295,12 +295,10 @@ class E2eKeysHandler(object):
         opentracing.log_kv({"one_time_keys": json_result, "failures": failures})
         return {"one_time_keys": json_result, "failures": failures}
 
-    @opentracing.trace_defered_function
     @defer.inlineCallbacks
     def upload_keys_for_user(self, user_id, device_id, keys):
         opentracing.set_tag("user_id", user_id)
         opentracing.set_tag("device_id", device_id)
-        opentracing.set_tag("keys", keys)
 
         time_now = self.clock.time_msec()
 
@@ -314,6 +312,13 @@ class E2eKeysHandler(object):
                 user_id,
                 time_now,
             )
+            opentracing.log_kv(
+                {
+                    "message": "Updating device_keys for user.",
+                    "user_id": user_id,
+                    "device_id": device_id,
+                }
+            )
             # TODO: Sign the JSON with the server key
             changed = yield self.store.set_e2e_device_keys(
                 user_id, device_id, time_now, device_keys
@@ -321,15 +326,26 @@ class E2eKeysHandler(object):
             if changed:
                 # Only notify about device updates *if* the keys actually changed
                 yield self.device_handler.notify_device_update(user_id, [device_id])
-
+        else:
+            opentracing.log_kv(
+                {"message": "Not updating device_keys for user", "user_id": user_id}
+            )
         one_time_keys = keys.get("one_time_keys", None)
+        opentracing.set_tag("one_time_keys", one_time_keys)
         if one_time_keys:
+            opentracing.log_kv(
+                {
+                    "message": "Updating one_time_keys for device.",
+                    "user_id": user_id,
+                    "device_id": device_id,
+                }
+            )
             yield self._upload_one_time_keys_for_user(
                 user_id, device_id, time_now, one_time_keys
             )
         else:
             opentracing.log_kv(
-                {"event": "did not upload one_time_keys", "reason": "no keys given"}
+                {"message": "Did not update one_time_keys", "reason": "no keys given"}
             )
 
         # the device should have been registered already, but it may have been
@@ -344,7 +360,6 @@ class E2eKeysHandler(object):
         opentracing.set_tag("one_time_key_counts", result)
         return {"one_time_key_counts": result}
 
-    @opentracing.trace_defered_function
     @defer.inlineCallbacks
     def _upload_one_time_keys_for_user(
         self, user_id, device_id, time_now, one_time_keys
@@ -386,6 +401,9 @@ class E2eKeysHandler(object):
                     (algorithm, key_id, encode_canonical_json(key).decode("ascii"))
                 )
 
+        opentracing.log_kv(
+            {"message": "Inserting new one_time_keys.", "keys": new_keys}
+        )
         yield self.store.add_e2e_one_time_keys(user_id, device_id, time_now, new_keys)
 
 
diff --git a/synapse/storage/end_to_end_keys.py b/synapse/storage/end_to_end_keys.py
index 4fa694e06a..29d91dfa6b 100644
--- a/synapse/storage/end_to_end_keys.py
+++ b/synapse/storage/end_to_end_keys.py
@@ -115,7 +115,6 @@ class EndToEndKeyWorkerStore(SQLBaseStore):
         opentracing.log_kv(result)
         return result
 
-    @opentracing.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
@@ -131,10 +130,6 @@ class EndToEndKeyWorkerStore(SQLBaseStore):
             key_id) to json string for key
         """
 
-        opentracing.set_tag("user_id", user_id)
-        opentracing.set_tag("device_id", device_id)
-        opentracing.set_tag("key_ids", key_ids)
-
         rows = yield self._simple_select_many_batch(
             table="e2e_one_time_keys_json",
             column="key_id",
@@ -143,8 +138,11 @@ class EndToEndKeyWorkerStore(SQLBaseStore):
             keyvalues={"user_id": user_id, "device_id": device_id},
             desc="add_e2e_one_time_keys_check",
         )
-
-        return {(row["algorithm"], row["key_id"]): row["key_json"] for row in rows}
+        result = {(row["algorithm"], row["key_id"]): row["key_json"] for row in rows}
+        opentracing.log_kv(
+            {"message": "Fetched one time keys for user", "one_time_keys": result}
+        )
+        return result
 
     @defer.inlineCallbacks
     def add_e2e_one_time_keys(self, user_id, device_id, time_now, new_keys):
@@ -159,7 +157,6 @@ class EndToEndKeyWorkerStore(SQLBaseStore):
                 (algorithm, key_id, key json)
         """
 
-        @opentracing.trace_function
         def _add_e2e_one_time_keys(txn):
             opentracing.set_tag("user_id", user_id)
             opentracing.set_tag("device_id", device_id)
@@ -219,7 +216,6 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
         or the keys were already in the database.
         """
 
-        @opentracing.trace_function
         def _set_e2e_device_keys_txn(txn):
             opentracing.set_tag("user_id", user_id)
             opentracing.set_tag("device_id", device_id)
@@ -239,7 +235,7 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
             new_key_json = encode_canonical_json(device_keys).decode("utf-8")
 
             if old_key_json == new_key_json:
-                opentracing.log_kv({"event", "key already stored"})
+                opentracing.log_kv({"Message", "Device key already stored."})
                 return False
 
             self._simple_upsert_txn(
@@ -248,7 +244,7 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
                 keyvalues={"user_id": user_id, "device_id": device_id},
                 values={"ts_added_ms": time_now, "key_json": new_key_json},
             )
-
+            opentracing.log_kv({"message": "Device keys stored."})
             return True
 
         return self.runInteraction("set_e2e_device_keys", _set_e2e_device_keys_txn)