Add advisory_lock utility from AWX

Author: AlanCoding

ansible_base/lib/utils/db.py

@@ -1,6 +1,7 @@
 from contextlib import contextmanager
+from zlib import crc32
 
-from django.db import connection, transaction
+from django.db import DEFAULT_DB_ALIAS, connection, connections, transaction
 from django.db.migrations.executor import MigrationExecutor
 
 
@@ -25,3 +26,114 @@ def migrations_are_complete() -> bool:
     executor = MigrationExecutor(connection)
     plan = executor.migration_plan(executor.loader.graph.leaf_nodes())
     return not bool(plan)
+
+
+# NOTE: the django_pglocks_advisory_lock context manager was forked from the django-pglocks v1.0.4
+# that was licensed under the MIT license
+
+
+@contextmanager
+def django_pglocks_advisory_lock(lock_id, shared=False, wait=True, using=None):
+
+    if using is None:
+        using = DEFAULT_DB_ALIAS
+
+    # Assemble the function name based on the options.
+
+    function_name = 'pg_'
+
+    if not wait:
+        function_name += 'try_'
+
+    function_name += 'advisory_lock'
+
+    if shared:
+        function_name += '_shared'
+
+    release_function_name = 'pg_advisory_unlock'
+    if shared:
+        release_function_name += '_shared'
+
+    # Format up the parameters.
+
+    tuple_format = False
+
+    if isinstance(
+        lock_id,
+        (
+            list,
+            tuple,
+        ),
+    ):
+        if len(lock_id) != 2:
+            raise ValueError("Tuples and lists as lock IDs must have exactly two entries.")
+
+        if not isinstance(lock_id[0], int) or not isinstance(lock_id[1], int):
+            raise ValueError("Both members of a tuple/list lock ID must be integers")
+
+        tuple_format = True
+    elif isinstance(lock_id, str):
+        # Generates an id within postgres integer range (-2^31 to 2^31 - 1).
+        # crc32 generates an unsigned integer in Py3, we convert it into
+        # a signed integer using 2's complement (this is a noop in Py2)
+        pos = crc32(lock_id.encode("utf-8"))
+        lock_id = (2**31 - 1) & pos
+        if pos & 2**31:
+            lock_id -= 2**31
+    elif not isinstance(lock_id, int):
+        raise ValueError("Cannot use %s as a lock id" % lock_id)
+
+    if tuple_format:
+        base = "SELECT %s(%d, %d)"
+        params = (
+            lock_id[0],
+            lock_id[1],
+        )
+    else:
+        base = "SELECT %s(%d)"
+        params = (lock_id,)
+
+    acquire_params = (function_name,) + params
+
+    command = base % acquire_params
+    cursor = connections[using].cursor()
+
+    cursor.execute(command)
+
+    if not wait:
+        acquired = cursor.fetchone()[0]
+    else:
+        acquired = True
+
+    try:
+        yield acquired
+    finally:
+        if acquired:
+            release_params = (release_function_name,) + params
+
+            command = base % release_params
+            cursor.execute(command)
+
+        cursor.close()
+
+
+@contextmanager
+def advisory_lock(*args, lock_session_timeout_milliseconds=0, **kwargs):
+    if connection.vendor == "postgresql":
+        cur = None
+        idle_in_transaction_session_timeout = None
+        idle_session_timeout = None
+        if lock_session_timeout_milliseconds > 0:
+            with connection.cursor() as cur:
+                idle_in_transaction_session_timeout = cur.execute("SHOW idle_in_transaction_session_timeout").fetchone()[0]
+                idle_session_timeout = cur.execute("SHOW idle_session_timeout").fetchone()[0]
+                cur.execute(f"SET idle_in_transaction_session_timeout = '{lock_session_timeout_milliseconds}'")
+                cur.execute(f"SET idle_session_timeout = '{lock_session_timeout_milliseconds}'")
+        with django_pglocks_advisory_lock(*args, **kwargs) as internal_lock:
+            yield internal_lock
+            if lock_session_timeout_milliseconds > 0:
+                with connection.cursor() as cur:
+                    cur.execute(f"SET idle_in_transaction_session_timeout = '{idle_in_transaction_session_timeout}'")
+                    cur.execute(f"SET idle_session_timeout = '{idle_session_timeout}'")
+    else:
+        yield True

docs/lib/advisory_lock.md

@@ -0,0 +1,35 @@
+## Database Named Locks
+
+Django-ansible-base hosts its own specialized utility for obtaining named locks.
+This follows the same contract as documented in the django-pglocks library
+
+https://pypi.org/project/django-pglocks/
+
+Due to a multitude of needs relevant to production use, discovered through its
+use in AWX, a number of points of divergence have emerged such as:
+
+ - the need to have it not error when running sqlite3 tests
+ - stuck processes holding the lock forever (adding pg-level idle timeout)
+
+The use for the purpose of a task would typically look like this
+
+```python
+from ansible_base.lib.utils.db import advisory_lock
+
+
+def my_task():
+    with advisory_lock('my_task_lock', wait=False) as held:
+        if held is False:
+            return
+        # continue to run logic in my_task
+```
+
+This is very useful to assure that no other process _in the cluster_ connected
+to the same postgres instance runs `my_task` at the same time as the process
+calling it here.
+
+The specific choice of `wait=False` and what to do when another task holds the lock,
+is the choice of the programmer in the specific case.
+In this case, the `return` would be okay in the situation where `my_task` is idempotent,
+and there is a "fallback" schedule in case a call was missed.
+The blocking/non-blocking choices are very dependent on the specific design and situation.

test_app/tests/lib/utils/test_db.py

@@ -1,9 +1,39 @@
+import threading
+import time
+
 import pytest
 
-from ansible_base.lib.utils.db import migrations_are_complete
+from ansible_base.lib.utils.db import advisory_lock, migrations_are_complete
 
 
 @pytest.mark.django_db
 def test_migrations_are_complete():
     "If you are running tests, migrations (test database) should be complete"
     assert migrations_are_complete()
+
+
+@pytest.mark.django_db
+def test_get_unclaimed_lock():
+    with advisory_lock('test_get_unclaimed_lock'):
+        pass
+
+
+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)
+
+
+@pytest.mark.django_db
+def test_determine_lock_is_held(django_db_blocker):
+    thread = threading.Thread(target=background_task, args=(django_db_blocker,))
+    thread.start()
+    for _ in range(5):
+        with advisory_lock('background_task_lock', wait=False) as held:
+            if held is False:
+                break
+        time.sleep(0.01)
+    else:
+        raise RuntimeError('Other thread never obtained lock')
+    thread.join()