Add advisory_lock utility from AWX (ansible#661)

This migrates `advisory_lock` from AWX to DAB. This is similar in spirit
to ansible#660

Importantly, I went the route of _not_ adding a new dependency, because
there was only a single method from the library we were using, and I'm
pretty confident that the MIT license lets you do whatever you want. In
doing this, I deleted the use of `six` for python2 compatibility which
further reduced the complexity of the code I ported.

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,126 @@ 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):
+    """Context manager that wraps the pglocks advisory lock
+
+    This obtains a named lock in postgres, idenfied by the args passed in
+    usually the lock identifier is a simple string.
+
+    @param: wait If True, block until the lock is obtained
+    @param: shared Whether or not the lock is shared
+    @param: lock_session_timeout_milliseconds Postgres-level timeout
+    @param: using django database identifier
+    """
+    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("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:
+                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:
+        raise RuntimeError(f'Advisory lock not implemented for database type {connection.vendor}')

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,65 @@
+import threading
+import time
+
 import pytest
+from django.db import connection
+from django.db.utils import OperationalError
 
-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()
+
+
+class TestAdvisoryLock:
+    @pytest.fixture(autouse=True)
+    def skip_if_sqlite(self):
+        if connection.vendor == 'sqlite':
+            pytest.skip('Advisory lock is not written for sqlite')
+
+    @pytest.mark.django_db
+    def test_get_unclaimed_lock(self):
+        with advisory_lock('test_get_unclaimed_lock'):
+            pass
+
+    @staticmethod
+    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(self, django_db_blocker):
+        thread = threading.Thread(target=TestAdvisoryLock.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()
+
+    @pytest.mark.django_db
+    def test_tuple_lock(self):
+        with advisory_lock([1234, 4321]):
+            pass
+
+    @pytest.mark.django_db
+    def test_invalid_tuple_name(self):
+        with pytest.raises(ValueError):
+            with advisory_lock(['test_invalid_tuple_name', 'foo']):
+                pass
+
+    @pytest.mark.django_db
+    def test_lock_session_timeout_milliseconds(self):
+        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)