Fix incomplete error handling in advisory_lock wrapper

Author: AlanCoding

ansible_base/lib/utils/db.py

@@ -1,10 +1,14 @@
 from contextlib import contextmanager
 from zlib import crc32
+import logging
 
-from django.db import DEFAULT_DB_ALIAS, connection, connections, transaction
+from django.db import DEFAULT_DB_ALIAS, connection, connections, transaction, OperationalError
 from django.db.migrations.executor import MigrationExecutor
 
 
+logger = logging.getLogger(__name__)
+
+
 @contextmanager
 def ensure_transaction():
     needs_new_transaction = not transaction.get_connection().in_atomic_block
@@ -129,6 +133,7 @@ def advisory_lock(*args, lock_session_timeout_milliseconds=0, **kwargs):
     @param: lock_session_timeout_milliseconds Postgres-level timeout
     @param: using django database identifier
     """
+    internal_error = False
     if connection.vendor == "postgresql":
         cur = None
         idle_in_transaction_session_timeout = None
@@ -139,12 +144,23 @@ def advisory_lock(*args, lock_session_timeout_milliseconds=0, **kwargs):
                 idle_session_timeout = cur.execute("SHOW idle_session_timeout").fetchone()[0]
                 cur.execute("SET idle_in_transaction_session_timeout = %s", (lock_session_timeout_milliseconds,))
                 cur.execute("SET idle_session_timeout = %s", (lock_session_timeout_milliseconds,))
-        with django_pglocks_advisory_lock(*args, **kwargs) as internal_lock:
-            yield internal_lock
-            if lock_session_timeout_milliseconds > 0:
+
+        try:
+            with django_pglocks_advisory_lock(*args, **kwargs) as internal_lock:
+                yield internal_lock
+        except OperationalError:
+            # Suspected case is that timeout happened due to the given timeout
+            # this is _expected_ to leave the connection in an unusable state, so dropping it is better
+            logger.info('Dropping connection due to suspected timeout inside advisory_lock')
+            connection.close_if_unusable_or_obsolete()
+            internal_error = True
+            raise
+        finally:
+            if (not internal_error) and lock_session_timeout_milliseconds > 0:
                 with connection.cursor() as cur:
                     cur.execute("SET idle_in_transaction_session_timeout = %s", (idle_in_transaction_session_timeout,))
                     cur.execute("SET idle_session_timeout = %s", (idle_session_timeout,))
+
     elif connection.vendor == "sqlite":
         yield True
     else:

test_app/tests/lib/utils/test_db.py

@@ -5,7 +5,8 @@
 from django.db import connection
 from django.db.utils import OperationalError
 
-from ansible_base.lib.utils.db import advisory_lock, migrations_are_complete
+from ansible_base.lib.utils.db import advisory_lock, migrations_are_complete, transaction
+from django.db import connection
 
 
 @pytest.mark.django_db
@@ -15,6 +16,8 @@ def test_migrations_are_complete():
 
 
 class TestAdvisoryLock:
+    THREAD_WAIT_TIME = 0.1
+
     @pytest.fixture(autouse=True)
     def skip_if_sqlite(self):
         if connection.vendor == 'sqlite':
@@ -30,7 +33,7 @@ def background_task(django_db_blocker):
         # HACK: as a thread the pytest.mark.django_db will not work
         django_db_blocker.unblock()
         with advisory_lock('background_task_lock'):
-            time.sleep(0.1)
+            time.sleep(TestAdvisoryLock.THREAD_WAIT_TIME)
 
     @pytest.mark.django_db
     def test_determine_lock_is_held(self, django_db_blocker):
@@ -40,7 +43,7 @@ def test_determine_lock_is_held(self, django_db_blocker):
             with advisory_lock('background_task_lock', wait=False) as held:
                 if held is False:
                     break
-            time.sleep(0.01)
+            time.sleep(TestAdvisoryLock.THREAD_WAIT_TIME / 5.0)
         else:
             raise RuntimeError('Other thread never obtained lock')
         thread.join()
@@ -56,10 +59,52 @@ def test_invalid_tuple_name(self):
             with advisory_lock(['test_invalid_tuple_name', 'foo']):
                 pass
 
-    @pytest.mark.django_db
-    def test_lock_session_timeout_milliseconds(self):
+
+# Special transaction=True parameter is used, because we do not want in normal test transactions
+# because dropping the connection would break the transaction context, erroring at end of test
+@pytest.mark.django_db(transaction=True)
+class TestAdvisoryLockPostgresErrors:
+    """Tests related to connection management and the advisory_lock"""
+
+    def kick_connection(self):
+        """
+        These (somewhat evil) tests are not gentle with the connection
+
+        At the start of one tests there is a good chance the connection is broken by the last test.
+        """
+        connection.ensure_connection()
+        # sanity checks
+        assert connection.connection
+        assert connection.connection.closed is False
+
+    def test_timeout_under_lock_by_sleep(self):
+        self.kick_connection()
+
         with pytest.raises(OperationalError) as exc:
             # uses miliseconds units
-            with advisory_lock('test_lock_session_timeout_milliseconds', lock_session_timeout_milliseconds=2):
-                time.sleep(3)
-        assert 'the connection is lost' in str(exc)
+            with advisory_lock('test_lock_session_timeout_milliseconds', lock_session_timeout_milliseconds=5):
+                time.sleep(0.1)
+
+        # This exception comes from the __exit__, either releasing the lock or closing the cursor
+        assert 'terminating connection due to idle-session timeout' in str(exc)
+
+        with connection.cursor() as cursor:
+            cursor.execute('SELECT 1')
+
+    def test_idle_after_exception(self):
+        self.kick_connection()
+
+        with pytest.raises(RuntimeError) as exc:
+            with advisory_lock('test_timeout_under_lock_with_query', lock_session_timeout_milliseconds=5):
+                with connection.cursor() as cursor:
+                    raise RuntimeError('exception from test')
+
+        assert 'exception from test' in str(exc)
+
+        time.sleep(0.1)
+
+        # The fact that this works shows that the timeout was reset in the context manager __exit__
+        # Prior bug was giving exception with
+        # consuming input failed: terminating connection due to idle-session timeout server closed the connection unexpectedly
+        with connection.cursor() as cursor:
+            cursor.execute('SELECT 1')