[fix] race condition

usecase/sandbox.py

@@ -3,7 +3,13 @@ from datetime import timedelta
 from uuid import uuid4
 
 from domain.sandbox import SandboxSession
-from usecase.interface import Clock, Logger, SandboxRuntime, SandboxSessionRepository
+from usecase.interface import (
+    Clock,
+    Logger,
+    SandboxLifecycleLocker,
+    SandboxRuntime,
+    SandboxSessionRepository,
+)
 
 
 @dataclass(frozen=True, slots=True)
@@ -15,93 +21,112 @@ class CreateSandbox:
     def __init__(
         self,
         repository: SandboxSessionRepository,
+        locker: SandboxLifecycleLocker,
         runtime: SandboxRuntime,
         clock: Clock,
         logger: Logger,
         ttl: timedelta,
     ) -> None:
         self._repository = repository
+        self._locker = locker
         self._runtime = runtime
         self._clock = clock
         self._logger = logger
         self._ttl = ttl
 
     def execute(self, command: CreateSandboxCommand) -> SandboxSession:
-        now = self._clock.now()
-        session = self._repository.get_active_by_chat_id(command.chat_id)
+        with self._locker.lock(command.chat_id):
+            session = self._repository.get_active_by_chat_id(command.chat_id)
+            now = self._clock.now()
 
-        if session is not None and session.expires_at > now:
+            if session is not None and session.expires_at > now:
+                self._logger.info(
+                    'sandbox_reused',
+                    attrs={
+                        'chat_id': command.chat_id,
+                        'session_id': session.session_id,
+                        'container_id': session.container_id,
+                    },
+                )
+                return session
+
+            if session is not None:
+                self._logger.info(
+                    'sandbox_replaced',
+                    attrs={
+                        'chat_id': command.chat_id,
+                        'session_id': session.session_id,
+                        'container_id': session.container_id,
+                    },
+                )
+                self._runtime.stop(session.container_id)
+                self._repository.delete(session.session_id)
+
+            created_at = self._clock.now()
+            expires_at = created_at + self._ttl
+            new_session = self._runtime.create(
+                session_id=_new_session_id(),
+                chat_id=command.chat_id,
+                created_at=created_at,
+                expires_at=expires_at,
+            )
+            self._repository.save(new_session)
             self._logger.info(
-                'sandbox_reused',
+                'sandbox_created',
                 attrs={
                     'chat_id': command.chat_id,
-                    'session_id': session.session_id,
-                    'container_id': session.container_id,
+                    'session_id': new_session.session_id,
+                    'container_id': new_session.container_id,
                 },
             )
-            return session
-
-        if session is not None:
-            self._logger.info(
-                'sandbox_replaced',
-                attrs={
-                    'chat_id': command.chat_id,
-                    'session_id': session.session_id,
-                    'container_id': session.container_id,
-                },
-            )
-            self._runtime.stop(session.container_id)
-            self._repository.delete(session.session_id)
-
-        expires_at = now + self._ttl
-        new_session = self._runtime.create(
-            session_id=_new_session_id(),
-            chat_id=command.chat_id,
-            created_at=now,
-            expires_at=expires_at,
-        )
-        self._repository.save(new_session)
-        self._logger.info(
-            'sandbox_created',
-            attrs={
-                'chat_id': command.chat_id,
-                'session_id': new_session.session_id,
-                'container_id': new_session.container_id,
-            },
-        )
-        return new_session
+            return new_session
 
 
 class CleanupExpiredSandboxes:
     def __init__(
         self,
         repository: SandboxSessionRepository,
+        locker: SandboxLifecycleLocker,
         runtime: SandboxRuntime,
         clock: Clock,
         logger: Logger,
     ) -> None:
         self._repository = repository
+        self._locker = locker
         self._runtime = runtime
         self._clock = clock
         self._logger = logger
 
     def execute(self) -> list[SandboxSession]:
-        now = self._clock.now()
-        expired_sessions = self._repository.list_expired(now)
+        expired_sessions = self._repository.list_expired(self._clock.now())
         cleaned_sessions: list[SandboxSession] = []
 
         for session in expired_sessions:
-            self._runtime.stop(session.container_id)
-            self._repository.delete(session.session_id)
-            cleaned_sessions.append(session)
-            self._logger.info(
-                'sandbox_cleaned',
-                attrs={
-                    'chat_id': session.chat_id,
-                    'session_id': session.session_id,
-                    'container_id': session.container_id,
-                },
-            )
+            with self._locker.lock(session.chat_id):
+                current_session = self._repository.get_active_by_chat_id(
+                    session.chat_id
+                )
+                now = self._clock.now()
+                if current_session is None:
+                    continue
+
+                if current_session.session_id != session.session_id:
+                    continue
+
+                if current_session.expires_at > now:
+                    continue
+
+                self._runtime.stop(current_session.container_id)
+                self._repository.delete(current_session.session_id)
+                cleaned_sessions.append(current_session)
+                self._logger.info(
+                    'sandbox_cleaned',
+                    attrs={
+                        'chat_id': current_session.chat_id,
+                        'session_id': current_session.session_id,
+                        'container_id': current_session.container_id,
+                    },
+                )
 
         return cleaned_sessions