Add async lock feature with distributed support

Signed-off-by: jamshale <jamiehalebc@gmail.com>

Author: jamshale

acapy_agent/config/default_context.py

@@ -132,6 +132,7 @@ async def load_plugins(self, context: InjectionContext):
             "acapy_agent.vc.data_integrity",
             "acapy_agent.wallet",
             "acapy_agent.wallet.keys",
+            "acapy_agent.core.async_lock",
         ]
 
         did_management_plugins = [

acapy_agent/core/async_lock/__init__.py

@@ -0,0 +1,34 @@
+"""Setup for async lock service."""
+
+import json
+
+from ...__main__ import LOGGER
+from ...config.injection_context import InjectionContext
+from ...utils.base_storage import get_postgres_connection_uri
+from .async_lock import AsyncLock
+from .async_lock_postgres import PostgresAsyncLock
+from .async_lock_sqlite import SqliteAsyncLock
+
+
+async def setup(context: InjectionContext):
+    """Set up the async lock service."""
+
+    storage_type = context.settings.get("wallet.storage_type")
+
+    match storage_type:
+        case "postgres_storage":
+            LOGGER.info("Setting up and binding PostgresAsyncLock...")
+            connection_uri = get_postgres_connection_uri(
+                json.loads(context.settings["wallet.storage_creds"]),
+                json.loads(context.settings["wallet.storage_config"]),
+            )
+            await PostgresAsyncLock.create(connection_uri)
+            context.injector.bind_instance(AsyncLock, PostgresAsyncLock())
+
+        case "default" | None:
+            LOGGER.info("Setting up and binding SqliteAsyncLock...")
+            SqliteAsyncLock.create()
+            context.injector.bind_instance(AsyncLock, SqliteAsyncLock())
+
+        case _:
+            raise ValueError(f"Unsupported storage type for async lock: {storage_type}")

acapy_agent/core/async_lock/async_lock.py

@@ -0,0 +1,18 @@
+"""Abstract base class for an asynchronous lock service."""
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+
+class AsyncLock(ABC):
+    """Abstract base class for an asynchronous lock service."""
+
+    @abstractmethod
+    async def create(cls, connection_uri: Optional[str] = None):
+        """Create an instance of the lock service."""
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    @abstractmethod
+    async def lock(self, value: str, timeout: float = 10.0) -> bool:
+        """Acquire a lock with the given value."""
+        raise NotImplementedError("Subclasses must implement this method.")

acapy_agent/core/async_lock/async_lock_postgres.py

@@ -0,0 +1,91 @@
+"""Advisory Lock Service for PostgreSQL."""
+
+import asyncio
+import hashlib
+import logging
+import struct
+import time
+from contextlib import asynccontextmanager
+from typing import Optional
+
+import psycopg
+from psycopg.rows import tuple_row
+from psycopg.sql import SQL, Identifier
+
+from .async_lock import AsyncLock
+
+LOGGER = logging.getLogger(__name__)
+
+
+class PostgresAsyncLock(AsyncLock):
+    """A class to manage distributed advisory locks in PostgreSQL.
+
+    Note: This implementation uses PostgreSQL's advisory locks
+    to provide a distributed locking mechanism suitable for clustered environments.
+    It is designed to work with PostgreSQL databases and requires a valid connection URI
+    that it acquires from the provided storage configuration.
+    """
+
+    @classmethod
+    async def create(cls, connection_uri: Optional[str] = None):
+        """Create a Lock instance with a PostgreSQL connection."""
+        cls.connection_uri = connection_uri
+        await cls._create_db_if_not_exists(cls.connection_uri)
+
+    @staticmethod
+    async def _create_db_if_not_exists(base_uri: str, dbname="lock_db"):
+        async with await psycopg.AsyncConnection.connect(
+            base_uri, autocommit=True
+        ) as conn:
+            async with conn.cursor(row_factory=tuple_row) as cur:
+                await cur.execute(
+                    "SELECT 1 FROM pg_database WHERE datname = %s", (dbname,)
+                )
+                exists = await cur.fetchone()
+
+                if not exists:
+                    LOGGER.info(f"Creating database '{dbname}'...")
+                    await cur.execute(
+                        SQL("CREATE DATABASE {}").format(Identifier(dbname))
+                    )
+                else:
+                    LOGGER.debug(f"Database '{dbname}' already exists.")
+
+    def _make_pg_lock_key(self, value) -> int:
+        h = hashlib.sha256(str(value).encode()).digest()
+        return struct.unpack("q", h[:8])[0]  # signed 64-bit int
+
+    @asynccontextmanager
+    async def lock(self, lock_key: str, timeout: int = 10):
+        """Acquires a PostgreSQL advisory lock for the given key.
+
+        Times out after `timeout` seconds if the lock isn't available.
+        """
+        lock_key = self._make_pg_lock_key(lock_key)
+        async with await psycopg.AsyncConnection.connect(
+            self.connection_uri, autocommit=True
+        ) as conn:
+            async with conn.cursor() as cur:
+                start = time.time()
+                acquired = False
+
+                while time.time() - start < timeout:
+                    await cur.execute("SELECT pg_try_advisory_lock(%s);", (lock_key,))
+                    acquired = (await cur.fetchone())[0]
+
+                    if acquired:
+                        LOGGER.debug(f"[LOCK ACQUIRED] Key: {lock_key}")
+                        break
+
+                    await asyncio.sleep(1)
+
+                if not acquired:
+                    raise TimeoutError(
+                        f"Could not acquire advisory lock {lock_key} within {timeout} seconds."
+                    )
+
+                try:
+                    yield
+                finally:
+                    await cur.execute("SELECT pg_advisory_unlock(%s);", (lock_key,))
+                    LOGGER.debug(f"[LOCK RELEASED] Key: {lock_key}")

acapy_agent/core/async_lock/async_lock_sqlite.py

@@ -0,0 +1,65 @@
+"""A module for handling asynchronous locks using SQLite and portalocker."""
+
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Optional
+
+import portalocker
+
+from .async_lock import AsyncLock
+
+LOGGER = logging.getLogger(__name__)
+
+
+class SqliteAsyncLock(AsyncLock):
+    """A class to handle asynchronous locking using SQLite.
+
+    Note: This implementation uses file-based locks with portalocker.
+    It is suitable for single-node applications where SQLite is used as the storage
+    backend.
+
+    It does not support distributed locking across multiple nodes. Clustered
+    environments should use the currently supported postgres lock or
+    implement a custom distributed lock mechanism.
+    """
+
+    @classmethod
+    def create(cls, connection_uri: Optional[str] = None):
+        """Create a Lock instance with a SQLite connection."""
+        cls.lock_dir = Path("/tmp/sqlite_locks")
+        cls.lock_dir.mkdir(parents=True, exist_ok=True)
+
+    @staticmethod
+    def _lock_file_path(lock_dir: Path, lock_name: str) -> Path:
+        return lock_dir / f"{lock_name}.lock"
+
+    @asynccontextmanager
+    async def lock(self, lock_key: str, timeout: int = 10):
+        """Acquire a lock with the given key using SQLite and portalocker."""
+        loop = asyncio.get_running_loop()
+        file_path = self._lock_file_path(self.lock_dir, lock_key)
+
+        lock = portalocker.Lock(
+            file_path,
+            mode="a+",
+            flags=portalocker.LOCK_EX,  # No fail_when_locked
+        )
+
+        async def acquire():
+            await loop.run_in_executor(None, lock.acquire)
+
+        try:
+            await asyncio.wait_for(acquire(), timeout=timeout)
+            LOGGER.debug(f"[LOCK ACQUIRED] Key: {lock_key}")
+            yield
+        except asyncio.TimeoutError:
+            LOGGER.warning(f"Timeout acquiring lock: {lock_key}")
+            raise TimeoutError(f"Timeout acquiring lock: {lock_key}")
+        finally:
+            try:
+                await loop.run_in_executor(None, lock.release)
+                LOGGER.debug(f"[LOCK RELEASED] Key: {lock_key}")
+            except Exception:
+                LOGGER.warning(f"Failed to release lock: {lock_key}")

acapy_agent/utils/base_storage.py

@@ -0,0 +1,21 @@
+"""Base storage utility for connection URI generation."""
+
+import urllib
+
+
+def get_postgres_connection_uri(storage_creds: dict, storage_config: dict) -> str:
+    """Get the connection URI for the PostgreSQL database."""
+    uri = "postgresql://"
+    config_url = storage_config.get("url")
+    if not config_url:
+        raise ValueError("No 'url' provided for postgres store")
+    if "account" not in storage_creds:
+        raise ValueError("No 'account' provided for postgres store")
+    if "password" not in storage_creds:
+        raise ValueError("No 'password' provided for postgres store")
+    account = urllib.parse.quote(storage_creds["account"])
+    password = urllib.parse.quote(storage_creds["password"])
+    # FIXME parse the URL, check for parameters, remove postgres:// prefix, etc
+    # config url expected to be in the form "host:port"
+    uri += f"{account}:{password}@{config_url}/postgres"
+    return uri