ongoing changes

Author: sanderegg

packages/service-library/src/servicelib/project_lock.py

@@ -1,87 +1,51 @@
 import datetime
-import logging
-from asyncio.log import logger
-from collections.abc import AsyncIterator
-from contextlib import asynccontextmanager
-from typing import Final, TypeAlias
+import functools
+from collections.abc import Callable, Coroutine
+from typing import Any, Final, ParamSpec, TypeAlias, TypeVar
 
 import redis
 import redis.exceptions
 from models_library.projects import ProjectID
 from models_library.projects_access import Owner
 from models_library.projects_state import ProjectLocked, ProjectStatus
-from redis.asyncio.lock import Lock
-from servicelib.redis._client import RedisClientSDK
 
-from .background_task import periodic_task
-from .logging_utils import log_context
-
-_logger = logging.getLogger(__name__)
+from .redis import RedisClientSDK, exclusive
 
 PROJECT_REDIS_LOCK_KEY: str = "project_lock:{}"
 PROJECT_LOCK_TIMEOUT: Final[datetime.timedelta] = datetime.timedelta(seconds=10)
-ProjectLock = Lock
 
 ProjectLockError: TypeAlias = redis.exceptions.LockError
 
 
-async def _auto_extend_project_lock(project_lock: Lock) -> None:
-    # NOTE: the background task already catches anything that might raise here
-    await project_lock.reacquire()
+P = ParamSpec("P")
+R = TypeVar("R")
 
 
-@asynccontextmanager
-async def lock_project(
-    redis_client: RedisClientSDK,
+def with_locked_project(
+    redis_client: RedisClientSDK | Callable[..., RedisClientSDK],
     *,
     project_uuid: str | ProjectID,
     status: ProjectStatus,
     owner: Owner | None = None,
-) -> AsyncIterator[None]:
-    """Context manager to lock and unlock a project by user_id
-
-    Raises:
-        ProjectLockError: if project is already locked
-    """
-
-    redis_lock = redis_client.create_lock(
-        PROJECT_REDIS_LOCK_KEY.format(project_uuid),
-        ttl=PROJECT_LOCK_TIMEOUT,
-    )
-
-    try:
-        if not await redis_lock.acquire(
-            blocking=False,
-            token=ProjectLocked(
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, R]]], Callable[P, Coroutine[Any, Any, R]]
+]:
+    def _decorator(
+        func: Callable[P, Coroutine[Any, Any, R]],
+    ) -> Callable[P, Coroutine[Any, Any, R]]:
+        @exclusive(
+            redis_client,
+            lock_key=PROJECT_REDIS_LOCK_KEY.format(project_uuid),
+            lock_value=ProjectLocked(
                 value=True,
                 owner=owner,
                 status=status,
             ).model_dump_json(),
-        ):
-            msg = f"Lock for project {project_uuid!r} owner {owner!r} could not be acquired"
-            raise ProjectLockError(msg)
+        )
+        @functools.wraps(func)
+        async def _wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            return await func(*args, **kwargs)
 
-        with log_context(
-            _logger,
-            logging.DEBUG,
-            msg=f"with lock for {owner=}:{project_uuid=}:{status=}",
-        ):
-            async with periodic_task(
-                _auto_extend_project_lock,
-                interval=0.6 * PROJECT_LOCK_TIMEOUT,
-                task_name=f"{PROJECT_REDIS_LOCK_KEY.format(project_uuid)}_lock_auto_extend",
-                project_lock=redis_lock,
-            ):
-                yield
+        return _wrapper
 
-    finally:
-        # let's ensure we release that stuff
-        try:
-            if await redis_lock.owned():
-                await redis_lock.release()
-        except (redis.exceptions.LockError, redis.exceptions.LockNotOwnedError) as exc:
-            logger.warning(
-                "releasing %s unexpectedly raised an exception: %s",
-                f"{redis_lock=!r}",
-                f"{exc}",
-            )
+    return _decorator

services/efs-guardian/src/simcore_service_efs_guardian/services/background_tasks.py

@@ -5,7 +5,7 @@
 from models_library.projects import ProjectID
 from models_library.projects_state import ProjectStatus
 from servicelib.logging_utils import log_context
-from servicelib.project_lock import lock_project
+from servicelib.project_lock import with_locked_project
 from simcore_postgres_database.utils_projects import (
     DBProjectNotFoundError,
     ProjectsRepo,
@@ -18,6 +18,23 @@
 _logger = logging.getLogger(__name__)
 
 
+async def _remove_data_with_lock(app: FastAPI, project_id: ProjectID) -> None:
+    # Decorate a new function that will call the necessary coroutine
+    efs_manager: EfsManager = app.state.efs_manager
+
+    @with_locked_project(
+        get_redis_lock_client(app),
+        project_uuid=project_id,
+        status=ProjectStatus.MAINTAINING,
+    )
+    async def _remove():
+        # Call the actual coroutine function
+        await efs_manager.remove_project_efs_data(project_id)
+
+    # Execute the decorated function
+    await _remove()
+
+
 async def removal_policy_task(app: FastAPI) -> None:
     _logger.info("Removal policy task started")
 
@@ -56,9 +73,4 @@ async def removal_policy_task(app: FastAPI) -> None:
                 logging.INFO,
                 msg=f"Removing data for project {project_id} started, project last change date {_project_last_change_date}, efs removal policy task age limit timedelta {app_settings.EFS_REMOVAL_POLICY_TASK_AGE_LIMIT_TIMEDELTA}",
             ):
-                async with lock_project(
-                    get_redis_lock_client(app),
-                    project_uuid=project_id,
-                    status=ProjectStatus.MAINTAINING,
-                ):
-                    await efs_manager.remove_project_efs_data(project_id)
+                await _remove_data_with_lock(app, project_id)