Use native UPSERTs where possible

.coveragerc

@@ -1,11 +1,7 @@
 [run]
 branch = True
 parallel = True
-source = synapse
-
-[paths]
-source=
-   coverage
+include = synapse/*
 
 [report]
 precision = 2

.gitignore

@@ -25,9 +25,9 @@ homeserver*.pid
 *.tls.dh
 *.tls.key
 
-.coverage
-.coverage.*
-!.coverage.rc
+.coverage*
+coverage.*
+!.coveragerc
 htmlcov
 
 demo/*/*.db

changelog.d/4306.misc

@@ -0,0 +1 @@
+Synapse will now take advantage of native UPSERT functionality in PostgreSQL 9.5+.

synapse/storage/_base.py

@@ -192,6 +192,58 @@ def __init__(self, db_conn, hs):
 
         self.database_engine = hs.database_engine
 
+        # We force simple upserts by default, and then enable them once we have
+        # finished background updates (checked by `_check_safe_to_upsert`).
+        self._force_simple_upsert = True
+
+        if getattr(hs.config, "_enable_native_upserts", True):
+            # Check ASAP (and then later, every 1s) to see if we have finished
+            # background updates.
+            self._clock.call_later(0.0, self._check_safe_to_upsert)
+
+    @defer.inlineCallbacks
+    def _check_safe_to_upsert(self):
+        """
+        Is it safe to use native UPSERT?
+
+        If there are background updates, we will need to wait, as they may be
+        the addition of indexes that set the UNIQUE constraint that we require.
+
+        If the background updates have not completed, wait a second and check again.
+        """
+        completed = yield self.has_completed_background_updates()
+
+        if completed:
+            # Now that indexes are built, we are allowed to do native UPSERTs if
+            # the underlying database supports it.
+            self._force_simple_upsert = False
+        else:
+            self._clock.call_later(1.0, self._check_safe_to_upsert)
+
+    @defer.inlineCallbacks
+    def has_completed_background_updates(self):
+        """
+        Check if all the background updates have completed. This is safe to run
+        on the master as well as slaves, and will be overridden by
+        background_updates.BackgroundUpdateStore if we are the master and have a
+        subclass with it.
+
+        Returns: Deferred[bool]: True if all background updates have completed
+        """
+        # otherwise, check if there are updates to be run. This is important,
+        # as we may be running on a worker which doesn't perform the bg updates
+        # itself, but still wants to wait for them to happen.
+        updates = yield self._simple_select_onecol(
+            "background_updates",
+            keyvalues=None,
+            retcol="1",
+            desc="check_background_updates",
+        )
+        if not updates:
+            defer.returnValue(True)
+
+        defer.returnValue(False)
+
     def start_profiling(self):
         self._previous_loop_ts = self._clock.time_msec()
 
@@ -494,8 +546,16 @@ def _simple_insert_many_txn(txn, table, values):
         txn.executemany(sql, vals)
 
     @defer.inlineCallbacks
-    def _simple_upsert(self, table, keyvalues, values,
-                       insertion_values={}, desc="_simple_upsert", lock=True):
+    def _simple_upsert(
+        self,
+        table,
+        keyvalues,
+        values,
+        insertion_values={},
+        desc="_simple_upsert",
+        lock=True,
+        best_effort=False,
+    ):
         """
 
         `lock` should generally be set to True (the default), but can be set
@@ -515,6 +575,7 @@ def _simple_upsert(self, table, keyvalues, values,
             insertion_values (dict): additional key/values to use only when
                 inserting
             lock (bool): True to lock the table when doing the upsert.
+            best_effort (bool): If we run into a transaction error, do we stop trying?
         Returns:
             Deferred(bool): True if a new entry was created, False if an
                 existing one was updated.
@@ -524,8 +585,13 @@ def _simple_upsert(self, table, keyvalues, values,
             try:
                 result = yield self.runInteraction(
                     desc,
-                    self._simple_upsert_txn, table, keyvalues, values, insertion_values,
-                    lock=lock
+                    self._simple_upsert_txn,
+                    table,
+                    keyvalues,
+                    values,
+                    insertion_values,
+                    lock=lock,
+                    best_effort=best_effort,
                 )
                 defer.returnValue(result)
             except self.database_engine.module.IntegrityError as e:
@@ -537,12 +603,45 @@ def _simple_upsert(self, table, keyvalues, values,
 
                 # presumably we raced with another transaction: let's retry.
                 logger.warn(
-                    "IntegrityError when upserting into %s; retrying: %s",
-                    table, e
+                    "%s when upserting into %s; retrying: %s", e.__name__, table, e
                 )
 
-    def _simple_upsert_txn(self, txn, table, keyvalues, values, insertion_values={},
-                           lock=True):
+    def _simple_upsert_txn(
+        self,
+        txn,
+        table,
+        keyvalues,
+        values,
+        insertion_values={},
+        lock=True,
+        best_effort=False,
+    ):
+        """
+        Pick the UPSERT method which works best on the platform. Either the
+        native one (Pg9.5+, recent SQLites), or fall back to an emulated method.
+        """
+        if self.database_engine.can_native_upsert and not self._force_simple_upsert:
+            return self._simple_upsert_txn_native_upsert(
+                txn,
+                table,
+                keyvalues,
+                values,
+                insertion_values=insertion_values,
+                best_effort=best_effort,
+            )
+        else:
+            return self._simple_upsert_txn_emulated(
+                txn,
+                table,
+                keyvalues,
+                values,
+                insertion_values=insertion_values,
+                lock=lock,
+            )
+
+    def _simple_upsert_txn_emulated(
+        self, txn, table, keyvalues, values, insertion_values={}, lock=True
+    ):
         # We need to lock the table :(, unless we're *really* careful
         if lock:
             self.database_engine.lock_table(txn, table)
@@ -551,7 +650,7 @@ def _simple_upsert_txn(self, txn, table, keyvalues, values, insertion_values={},
         sql = "UPDATE %s SET %s WHERE %s" % (
             table,
             ", ".join("%s = ?" % (k,) for k in values),
-            " AND ".join("%s = ?" % (k,) for k in keyvalues)
+            " AND ".join("%s = ?" % (k,) for k in keyvalues),
         )
         sqlargs = list(values.values()) + list(keyvalues.values())
 
@@ -569,12 +668,56 @@ def _simple_upsert_txn(self, txn, table, keyvalues, values, insertion_values={},
         sql = "INSERT INTO %s (%s) VALUES (%s)" % (
             table,
             ", ".join(k for k in allvalues),
-            ", ".join("?" for _ in allvalues)
+            ", ".join("?" for _ in allvalues),
         )
         txn.execute(sql, list(allvalues.values()))
         # successfully inserted
         return True
 
+    def _simple_upsert_txn_native_upsert(
+        self, txn, table, keyvalues, values, insertion_values={}, best_effort=False
+    ):
+        """
+        Use the native UPSERT functionality in recent PostgreSQL versions.
+        """
+        allvalues = {}
+        allvalues.update(keyvalues)
+        allvalues.update(values)
+        allvalues.update(insertion_values)
+
+        sql = (
+            "INSERT INTO %s (%s) VALUES (%s) "
+            "ON CONFLICT (%s) DO UPDATE SET %s "
+            "RETURNING (xmax = 0) AS inserted"
+        ) % (
+            table,
+            ", ".join(k for k in allvalues),
+            ", ".join("?" for _ in allvalues),
+            ", ".join(k for k in keyvalues),
+            ", ".join(k + "=EXCLUDED." + k for k in values),
+        )
+        try:
+            txn.execute(sql, list(allvalues.values()))
+        except self.database_engine.module.OperationalError as e:
+            # We only care about serialization errors, so check for it
+            if e.args[0] == "could not serialize access due to concurrent update":
+                # A concurrent update problem is when we try and do a native
+                # UPSERT but the row has changed from under us. We can either
+                # retry, or give up if asked to do so.
+                if best_effort:
+                    # If it's a concurrent-update problem, and this is marked as
+                    # 'best effort' (i.e. if there's a race, then the one we
+                    # raced with will suffice), then pretend that we succeeded.
+                    return False
+            else:
+                # Otherwise, raise, because it's a real OperationalError and we
+                # will need to be rolled back and retried.
+                raise
+
+        # One-tuple, which is a boolean for insertion or not
+        res = txn.fetchone()
+        return res[0]
+
     def _simple_select_one(self, table, keyvalues, retcols,
                            allow_none=False, desc="_simple_select_one"):
         """Executes a SELECT query on the named table, which is expected to

synapse/storage/client_ips.py

@@ -65,6 +65,15 @@ def __init__(self, db_conn, hs):
             columns=["last_seen"],
         )
 
+        # Register a unique index
+        self.register_background_index_update(
+            "user_ips_device_unique_index",
+            index_name="user_ips_device_unique_id",
+            table="user_ips",
+            columns=["user_id", "access_token", "ip", "user_agent", "device_id"],
+            unique=True,
+        )
+
         # (user_id, access_token, ip) -> (user_agent, device_id, last_seen)
         self._batch_row_update = {}
 
@@ -114,7 +123,8 @@ def update():
         )
 
     def _update_client_ips_batch_txn(self, txn, to_update):
-        self.database_engine.lock_table(txn, "user_ips")
+        if self._force_simple_upsert or (not self.database_engine.can_native_upsert):
+            self.database_engine.lock_table(txn, "user_ips")
 
         for entry in iteritems(to_update):
             (user_id, access_token, ip), (user_agent, device_id, last_seen) = entry

synapse/storage/engines/__init__.py

@@ -18,7 +18,7 @@
 
 from ._base import IncorrectDatabaseSetup
 from .postgres import PostgresEngine
-from .sqlite3 import Sqlite3Engine
+from .sqlite import Sqlite3Engine
 
 SUPPORTED_MODULE = {
     "sqlite3": Sqlite3Engine,

synapse/storage/engines/postgres.py

@@ -38,6 +38,18 @@ def convert_param_style(self, sql):
         return sql.replace("?", "%s")
 
     def on_new_connection(self, db_conn):
+
+        # Get the version of PostgreSQL that we're using. As per the psycopg2
+        # docs: The number is formed by converting the major, minor, and
+        # revision numbers into two-decimal-digit numbers and appending them
+        # together. For example, version 8.1.5 will be returned as 80105.
+        server_version = str(db_conn.server_version)
+        self._version = (
+            int(server_version[:-4], 10),
+            int(server_version[-4:-2], 10),
+            int(server_version[-2:], 10)
+        )
+
         db_conn.set_isolation_level(
             self.module.extensions.ISOLATION_LEVEL_REPEATABLE_READ
         )
@@ -54,6 +66,13 @@ def on_new_connection(self, db_conn):
 
         cursor.close()
 
+    @property
+    def can_native_upsert(self):
+        """
+        Can we use native UPSERTs? This requires PostgreSQL 9.5+.
+        """
+        return self._version >= (9, 5, 0)
+
     def is_deadlock(self, error):
         if isinstance(error, self.module.DatabaseError):
             # https://www.postgresql.org/docs/current/static/errcodes-appendix.html

synapse/storage/engines/sqlite3.py → synapse/storage/engines/sqlite.py

@@ -30,6 +30,21 @@ def __init__(self, database_module, database_config):
         self._current_state_group_id = None
         self._current_state_group_id_lock = threading.Lock()
 
+    @property
+    def can_native_upsert(self):
+        """
+        Do we support native UPSERTs? This requires SQLite3 3.24+, plus some
+        more work we haven't done yet to tell what was inserted vs updated.
+        """
+        # SQLite currently doesn't have a way of telling if an UPSERT ended up
+        # with an INSERT or an UPDATE
+        # To enable, uncomment:
+        #
+        # from sqlite3 import sqlite_version_info
+        # return sqlite_version_info >= (3, 24, 0)
+
+        return False
+
     def check_database(self, txn):
         pass
 

synapse/storage/monthly_active_users.py

@@ -237,7 +237,9 @@ def upsert_monthly_active_user_txn(self, txn, user_id):
         # Am consciously deciding to lock the table on the basis that is ought
         # never be a big table and alternative approaches (batching multiple
         # upserts into a single txn) introduced a lot of extra complexity.
-        # See https://github.com/matrix-org/synapse/issues/3854 for more
+        # See https://github.com/matrix-org/synapse/issues/3854 for more.
+        # If we support native upserts, we'll not lock, but also not retry
+        # on any races, by setting best_effort=True.
         is_insert = self._simple_upsert_txn(
             txn,
             table="monthly_active_users",
@@ -247,6 +249,7 @@ def upsert_monthly_active_user_txn(self, txn, user_id):
             values={
                 "timestamp": int(self._clock.time_msec()),
             },
+            best_effort=True,
         )
 
         return is_insert

synapse/storage/prepare_database.py

@@ -25,7 +25,7 @@
 
 # Remember to update this number every time a change is made to database
 # schema files, so the users will be informed on server restarts.
-SCHEMA_VERSION = 53
+SCHEMA_VERSION = 54
 
 dir_path = os.path.abspath(os.path.dirname(__file__))
 

synapse/storage/schema/delta/54/user_ips_index.sql

@@ -0,0 +1,18 @@
+/* Copyright 2018 New Vector Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+-- add a new unique index to user_ips table
+INSERT INTO background_updates (update_name, progress_json) VALUES
+  ('user_ips_device_unique_index', '{}');