[v2] Reduce session database locking

.changes/next-release/enhancement-telemetry-83742.json

@@ -0,0 +1,5 @@
+{
+  "type": "enhancement",
+  "category": "``telemetry``",
+  "description": "Reduce session database locking which could cause slowdowns when the AWS CLI cache directory is on a network filesystem."
+}

awscli/telemetry.py

@@ -36,6 +36,12 @@
 _DATABASE_FILENAME = 'session.db'
 _SESSION_LENGTH_SECONDS = 60 * 30
 _SESSION_ID_LENGTH = 12
+# How stale the stored timestamp must be before it is written. Refreshing on every
+# invocation takes a write lock each time, which can cause lock contention.
+_TIMESTAMP_REFRESH_SECONDS = 60
+# How long to wait for a database lock before giving up. Session data is not
+# critical, so it's better to skip it than to make the user wait
+_BUSY_TIMEOUT_SECONDS = 0.1
 # Set to "true" to skip collecting session ids entirely. Opting out avoids
 # opening the session database, whose file locking can be expensive when the
 # database lives on a network filesystem.
@@ -90,6 +96,7 @@ def __init__(self, connection=None, cache_dir=None):
             self._cache_dir / _DATABASE_FILENAME,
             check_same_thread=False,
             isolation_level=None,
+            timeout=_BUSY_TIMEOUT_SECONDS,
         )
         self._ensure_database_setup()
 
@@ -196,6 +203,12 @@ def read_host_id(self):
 
 
 class CLISessionDatabaseSweeper:
+    _CHECK_EXPIRED = """
+        SELECT 1
+        FROM session
+        WHERE timestamp < ?
+        LIMIT 1
+    """
     _DELETE_RECORDS = """
         DELETE FROM session
         WHERE timestamp < ?
@@ -206,6 +219,17 @@ def __init__(self, connection):
 
     def sweep(self, timestamp):
         try:
+            # A DELETE takes a write lock even when it matches no rows, which
+            # is the common case since expired records are rare. Reads don't
+            # take a write lock, so check for expired records first and only
+            # issue the DELETE when there's something to remove. This matters
+            # on network filesystems (e.g. NFS/EFS) where each lock is a
+            # network round trip.
+            has_expired = self._connection.execute(
+                self._CHECK_EXPIRED, (timestamp,)
+            ).fetchone()
+            if has_expired is None:
+                return
             self._connection.execute(self._DELETE_RECORDS, (timestamp,))
         except Exception:
             # This is just a background cleanup task. No need to
@@ -254,14 +278,23 @@ def _session_id(self):
     def session_id(self):
         if (cached_data := self._reader.read(self.cache_key)) is not None:
             # Cache hit, but session id is expired. Generate new id and update.
-            if (
+            is_expired = (
                 cached_data.timestamp + _SESSION_LENGTH_SECONDS
                 < self._timestamp
-            ):
+            )
+            if is_expired:
                 cached_data.session_id = self._session_id
-            # Always update the timestamp to last used.
-            cached_data.timestamp = self._timestamp
-            self._writer.write(cached_data)
+            # Update the timestamp to last used, but only if the stored value
+            # is older than _TIMESTAMP_REFRESH_SECONDS. A session can expire
+            # up to this many seconds earlier than its last actual use, but
+            # avoids a write on every invocation if close together.
+            if (
+                is_expired
+                or self._timestamp - cached_data.timestamp
+                > _TIMESTAMP_REFRESH_SECONDS
+            ):
+                cached_data.timestamp = self._timestamp
+                self._writer.write(cached_data)
             return cached_data.session_id
         # Cache miss, generate and write new record.
         session_id = self._session_id

tests/functional/test_telemetry.py

@@ -20,6 +20,7 @@
 
 from awscli.clidriver import create_clidriver
 from awscli.telemetry import (
+    _TIMESTAMP_REFRESH_SECONDS,
     CLISessionData,
     CLISessionDatabaseConnection,
     CLISessionDatabaseReader,
@@ -197,6 +198,26 @@ def test_sweep_never_raises(self, session_sweeper):
         # but the `sweep` method catches bare exceptions.
         session_sweeper.sweep({'bad': 'input'})
 
+    def test_sweep_does_not_delete_when_nothing_expired(
+        self, expired_data, session_conn, session_sweeper
+    ):
+        with patch.object(
+            session_conn, 'execute', wraps=session_conn.execute
+        ) as spy:
+            session_sweeper.sweep(1000000000)
+        queries = [call.args[0] for call in spy.call_args_list]
+        assert not any('DELETE' in query for query in queries)
+
+    def test_sweep_deletes_when_expired(
+        self, expired_data, session_conn, session_sweeper
+    ):
+        with patch.object(
+            session_conn, 'execute', wraps=session_conn.execute
+        ) as spy:
+            session_sweeper.sweep(1000000001)
+        queries = [call.args[0] for call in spy.call_args_list]
+        assert any('DELETE' in query for query in queries)
+
 
 class TestCLISessionGenerator:
     def test_generate_session_id(self, session_generator):
@@ -307,8 +328,8 @@ def test_cached_session_id_not_updated_if_valid(
         session_data_1 = session_reader.read(orchestrator_1.cache_key)
         assert session_data_1.session_id == session_id_1
 
-        # Update the timestamp.
-        patched_time.return_value = 5555555556
+        # Advance past the refresh window so the timestamp is rewritten.
+        patched_time.return_value = 5555555555 + _TIMESTAMP_REFRESH_SECONDS + 1
         orchestrator_2 = CLISessionOrchestrator(
             session_generator, session_writer, session_reader, session_sweeper
         )
@@ -321,9 +342,47 @@ def test_cached_session_id_not_updated_if_valid(
         assert session_data_2.session_id == session_id_2
         assert session_data_2.session_id == session_data_1.session_id
         # Only timestamp should be updated.
-        assert session_data_2.timestamp == 5555555556
+        assert (
+            session_data_2.timestamp
+            == 5555555555 + _TIMESTAMP_REFRESH_SECONDS + 1
+        )
         assert session_data_2.timestamp != session_data_1.timestamp
 
+    def test_cached_timestamp_not_rewritten_within_refresh_window(
+        self,
+        patched_tty_name,
+        patched_time,
+        patched_stdin,
+        session_sweeper,
+        session_generator,
+        session_reader,
+        session_writer,
+    ):
+        # Rewriting the timestamp takes a write lock, so an invocation within
+        # the refresh window should leave the stored record untouched.
+        patched_stdin.fileno.return_value = None
+
+        orchestrator_1 = CLISessionOrchestrator(
+            session_generator, session_writer, session_reader, session_sweeper
+        )
+        session_id_1 = orchestrator_1.session_id
+        session_data_1 = session_reader.read(orchestrator_1.cache_key)
+
+        # Advance time, but stay inside the refresh window.
+        patched_time.return_value = 5555555555 + _TIMESTAMP_REFRESH_SECONDS - 1
+        orchestrator_2 = CLISessionOrchestrator(
+            session_generator, session_writer, session_reader, session_sweeper
+        )
+        session_id_2 = orchestrator_2.session_id
+        session_data_2 = session_reader.read(orchestrator_2.cache_key)
+
+        # The same session id is still returned to the caller.
+        assert session_id_2 == session_id_1
+        assert session_data_2.session_id == session_data_1.session_id
+        # But the stored timestamp was not rewritten.
+        assert session_data_2.timestamp == session_data_1.timestamp
+        assert session_data_2.timestamp == 5555555555
+
 
 def test_register_session_id_event_injects_sid_on_before_create_client():
     session = MagicMock(Session)