use stream ID generator instead of timestamp

3 files changed, 20 insertions, 19 deletions

diff --git a/synapse/storage/__init__.py b/synapse/storage/__init__.py
index 0a64f90624..e9a9c2cd8d 100644
--- a/synapse/storage/__init__.py
+++ b/synapse/storage/__init__.py
@@ -136,6 +136,9 @@ class DataStore(
         self._device_list_id_gen = StreamIdGenerator(
             db_conn, "device_lists_stream", "stream_id"
         )
+        self._cross_signing_id_gen = StreamIdGenerator(
+            db_conn, "e2e_cross_signing_keys", "stream_id"
+        )
 
         self._access_tokens_id_gen = IdGenerator(db_conn, "access_tokens", "id")
         self._event_reports_id_gen = IdGenerator(db_conn, "event_reports", "id")
diff --git a/synapse/storage/end_to_end_keys.py b/synapse/storage/end_to_end_keys.py
index b218b7b2e8..4b37bffb0b 100644
--- a/synapse/storage/end_to_end_keys.py
+++ b/synapse/storage/end_to_end_keys.py
@@ -282,7 +282,7 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
             "delete_e2e_keys_by_device", delete_e2e_keys_by_device_txn
         )
 
-    def _set_e2e_cross_signing_key_txn(self, txn, user_id, key_type, key, added_ts):
+    def _set_e2e_cross_signing_key_txn(self, txn, user_id, key_type, key):
         """Set a user's cross-signing key.
 
         Args:
@@ -292,7 +292,6 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
                 for a master key, 'self_signing' for a self-signing key, or
                 'user_signing' for a user-signing key
             key (dict): the key data
-            added_ts (int): the timestamp for when the key was added
         """
         # the cross-signing keys need to occupy the same namespace as devices,
         # since signatures are identified by device ID.  So add an entry to the
@@ -327,25 +326,25 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
         )
 
         # and finally, store the key itself
-        self._simple_insert(
-            "e2e_cross_signing_keys",
-            values={
-                "user_id": user_id,
-                "keytype": key_type,
-                "keydata": json.dumps(key),
-                "added_ts": added_ts,
-            },
-            desc="store_master_key",
-        )
+        with self._cross_signing_id_gen.get_next() as stream_id:
+            self._simple_insert(
+                "e2e_cross_signing_keys",
+                values={
+                    "user_id": user_id,
+                    "keytype": key_type,
+                    "keydata": json.dumps(key),
+                    "stream_id": stream_id,
+                },
+                desc="store_master_key",
+            )
 
-    def set_e2e_cross_signing_key(self, user_id, key_type, key, added_ts):
+    def set_e2e_cross_signing_key(self, user_id, key_type, key):
         """Set a user's cross-signing key.
 
         Args:
             user_id (str): the user to set the user-signing key for
             key_type (str): the type of cross-signing key to set
             key (dict): the key data
-            added_ts (int): the timestamp for when the key was added
         """
         return self.runInteraction(
             "add_e2e_cross_signing_key",
@@ -353,7 +352,6 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
             user_id,
             key_type,
             key,
-            added_ts,
         )
 
     def _get_e2e_cross_signing_key_txn(self, txn, user_id, key_type, from_user_id=None):
@@ -374,7 +372,7 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
         sql = (
             "SELECT keydata "
             "  FROM e2e_cross_signing_keys "
-            " WHERE user_id = ? AND keytype = ? ORDER BY added_ts DESC LIMIT 1"
+            " WHERE user_id = ? AND keytype = ? ORDER BY stream_id DESC LIMIT 1"
         )
         txn.execute(sql, (user_id, key_type))
         row = txn.fetchone()
diff --git a/synapse/storage/schema/delta/56/signing_keys.sql b/synapse/storage/schema/delta/56/signing_keys.sql
index 6a9ef1782e..27a96123e3 100644
--- a/synapse/storage/schema/delta/56/signing_keys.sql
+++ b/synapse/storage/schema/delta/56/signing_keys.sql
@@ -20,11 +20,11 @@ CREATE TABLE IF NOT EXISTS e2e_cross_signing_keys (
     keytype TEXT NOT NULL,
     -- the full key information, as a json-encoded dict
     keydata TEXT NOT NULL,
-    -- time that the key was added
-    added_ts BIGINT NOT NULL
+    -- for keeping the keys in order, so that we can fetch the latest one
+    stream_id BIGINT NOT NULL
 );
 
-CREATE UNIQUE INDEX e2e_cross_signing_keys_idx ON e2e_cross_signing_keys(user_id, keytype, added_ts);
+CREATE UNIQUE INDEX e2e_cross_signing_keys_idx ON e2e_cross_signing_keys(user_id, keytype, stream_id);
 
 -- cross-signing signatures
 CREATE TABLE IF NOT EXISTS e2e_cross_signing_signatures (