refactor

Author: sanderegg

services/web/server/src/simcore_service_webserver/garbage_collector/_core_disconnected.py

@@ -30,7 +30,7 @@ async def remove_disconnected_user_resources(
 
     # clean up all resources of expired keys
     for dead_session in dead_user_sessions:
-        user_id = dead_session["user_id"]
+        user_id = dead_session.user_id
 
         # (0) If key has no resources => remove from registry and continue
         resources = await registry.get_resources(dead_session)
@@ -63,7 +63,7 @@ async def remove_disconnected_user_resources(
                     await _projects_service.close_project_for_user(
                         user_id=user_id,
                         project_uuid=project_id,
-                        client_session_id=dead_session["client_session_id"],
+                        client_session_id=dead_session.client_session_id,
                         app=app,
                         simcore_user_agent=UNDEFINED_DEFAULT_SIMCORE_USER_AGENT_VALUE,
                         wait_for_service_closed=True,

services/web/server/src/simcore_service_webserver/garbage_collector/_core_guests.py

@@ -192,7 +192,7 @@ async def remove_users_manually_marked_as_guests(
     ) = await registry.get_all_resource_keys()
 
     skip_users = {
-        int(user_session["user_id"])
+        user_session.user_id
         for user_session in itertools.chain(
             all_user_session_alive, all_user_sessions_dead
         )

services/web/server/src/simcore_service_webserver/projects/_projects_service.py

@@ -110,9 +110,9 @@
     get_redis_document_manager_client_sdk,
     get_redis_lock_manager_client_sdk,
 )
+from ..resource_manager.models import UserSession
 from ..resource_manager.user_sessions import (
     PROJECT_ID_KEY,
-    UserSessionID,
     managed_resource,
 )
 from ..resource_usage import service as rut_api
@@ -1441,7 +1441,7 @@ async def post_trigger_connected_service_retrieve(
 
 
 async def _user_has_another_active_session(
-    users_sessions_ids: list[UserSessionID], app: web.Application
+    users_sessions_ids: list[UserSession], app: web.Application
 ) -> bool:
     # NOTE if there is an active socket in use, that means the client is active
     for u in users_sessions_ids:
@@ -1452,7 +1452,7 @@ async def _user_has_another_active_session(
 
 
 async def _clean_user_disconnected_clients(
-    users_sessions_ids: list[UserSessionID], app: web.Application
+    users_sessions_ids: list[UserSession], app: web.Application
 ):
     for u in users_sessions_ids:
         with managed_resource(u.user_id, u.client_session_id, app) as user_session:

services/web/server/src/simcore_service_webserver/resource_manager/models.py

@@ -0,0 +1,58 @@
+from typing import Final, Self, TypeAlias, TypedDict
+
+from models_library.basic_types import UUIDStr
+from models_library.users import UserID
+from pydantic import BaseModel, ConfigDict
+from pydantic.config import JsonDict
+
+ALIVE_SUFFIX: Final[str] = "alive"  # points to a string type
+RESOURCE_SUFFIX: Final[str] = "resources"  # points to a hash (like a dict) type
+RedisHashKey: TypeAlias = str
+
+
+class UserSession(BaseModel):
+    """Parts of the key used in redis for a user-session"""
+
+    user_id: UserID
+    client_session_id: str
+
+    def to_redis_hash_key(self) -> RedisHashKey:
+        return ":".join(f"{k}={v}" for k, v in self.model_dump().items())
+
+    @classmethod
+    def from_redis_hash_key(cls, hash_key: RedisHashKey) -> Self:
+        key = dict(x.split("=") for x in hash_key.split(":") if "=" in x)
+        return cls.model_validate(key)
+
+    @staticmethod
+    def _update_json_schema_extra(schema: JsonDict) -> None:
+        schema.update(
+            {
+                "exampels": [
+                    {
+                        "user_id": 7,
+                        "client_session_id": "c7fc4985-f96a-4be3-a8ed-5a43b1aa15e2",
+                    },
+                    {
+                        "user_id": 666,
+                        "client_session_id": "*",
+                    },
+                ]
+            }
+        )
+
+    model_config = ConfigDict(
+        frozen=True,
+        json_schema_extra=_update_json_schema_extra,
+    )
+
+
+class ResourcesDict(TypedDict, total=False):
+    """Field-value pairs of {user_id}:{client_session_id}:resources key"""
+
+    project_id: UUIDStr
+    socket_id: str
+
+
+AliveSessions: TypeAlias = list[UserSession]
+DeadSessions: TypeAlias = list[UserSession]

services/web/server/src/simcore_service_webserver/resource_manager/registry.py

@@ -14,19 +14,21 @@
 """
 
 import logging
-from typing import TypeAlias
 
 import redis.asyncio as aioredis
 from aiohttp import web
-from models_library.basic_types import UUIDStr
-from models_library.users import UserID
 from servicelib.redis import handle_redis_returns_union_types
-from typing_extensions import (  # https://docs.pydantic.dev/latest/api/standard_library_types/#typeddict
-    TypedDict,
-)
 
 from ..redis import get_redis_resources_client
 from ._constants import APP_CLIENT_SOCKET_REGISTRY_KEY
+from .models import (
+    ALIVE_SUFFIX,
+    RESOURCE_SUFFIX,
+    AliveSessions,
+    DeadSessions,
+    ResourcesDict,
+    UserSession,
+)
 
 _logger = logging.getLogger(__name__)
 
@@ -35,26 +37,6 @@
 #        Key: user_id=1:client_session_id=7f40353b-db02-4474-a44d-23ce6a6e428c:alive = 1
 #        Key: user_id=1:client_session_id=7f40353b-db02-4474-a44d-23ce6a6e428c:resources = {project_id: ... , socket_id: ...}
 #
-_ALIVE_SUFFIX = "alive"  # points to a string type
-_RESOURCE_SUFFIX = "resources"  # points to a hash (like a dict) type
-
-
-class UserSession(TypedDict):
-    """Parts of the key used in redis for a user-session"""
-
-    user_id: UserID
-    client_session_id: str
-
-
-class ResourcesDict(TypedDict, total=False):
-    """Field-value pairs of {user_id}:{client_session_id}:resources key"""
-
-    project_id: UUIDStr
-    socket_id: str
-
-
-AliveSessions: TypeAlias = list[UserSession]
-DeadSessions: TypeAlias = list[UserSession]
 
 
 class RedisResourceRegistry:
@@ -75,17 +57,12 @@ def __init__(self, app: web.Application):
     def app(self) -> web.Application:
         return self._app
 
-    @classmethod
-    def _hash_key(cls, key: UserSession) -> str:
-        hash_key: str = ":".join(f"{k}={v}" for k, v in key.items())
-        return hash_key
-
     @classmethod
     def _decode_hash_key(cls, hash_key: str) -> UserSession:
         tmp_key = (
-            hash_key[: -len(f":{_RESOURCE_SUFFIX}")]
-            if hash_key.endswith(f":{_RESOURCE_SUFFIX}")
-            else hash_key[: -len(f":{_ALIVE_SUFFIX}")]
+            hash_key[: -len(f":{RESOURCE_SUFFIX}")]
+            if hash_key.endswith(f":{RESOURCE_SUFFIX}")
+            else hash_key[: -len(f":{ALIVE_SUFFIX}")]
         )
         key = dict(x.split("=") for x in tmp_key.split(":"))
         return UserSession(**key)  # type: ignore
@@ -96,27 +73,27 @@ def client(self) -> aioredis.Redis:
         return client
 
     async def set_resource(self, key: UserSession, resource: tuple[str, str]) -> None:
-        hash_key = f"{self._hash_key(key)}:{_RESOURCE_SUFFIX}"
+        hash_key = f"{key.to_redis_hash_key()}:{RESOURCE_SUFFIX}"
         field, value = resource
         await handle_redis_returns_union_types(
             self.client.hset(hash_key, mapping={field: value})
         )
 
     async def get_resources(self, key: UserSession) -> ResourcesDict:
-        hash_key = f"{self._hash_key(key)}:{_RESOURCE_SUFFIX}"
+        hash_key = f"{key.to_redis_hash_key()}:{RESOURCE_SUFFIX}"
         fields = await handle_redis_returns_union_types(self.client.hgetall(hash_key))
         return ResourcesDict(**fields)
 
     async def remove_resource(self, key: UserSession, resource_name: str) -> None:
-        hash_key = f"{self._hash_key(key)}:{_RESOURCE_SUFFIX}"
+        hash_key = f"{key.to_redis_hash_key()}:{RESOURCE_SUFFIX}"
         await handle_redis_returns_union_types(
             self.client.hdel(hash_key, resource_name)
         )
 
     async def find_resources(self, key: UserSession, resource_name: str) -> list[str]:
         resources: list[str] = []
         # the key might only be partialy complete
-        partial_hash_key = f"{self._hash_key(key)}:{_RESOURCE_SUFFIX}"
+        partial_hash_key = f"{key.to_redis_hash_key()}:{RESOURCE_SUFFIX}"
         async for scanned_key in self.client.scan_iter(match=partial_hash_key):
             if await handle_redis_returns_union_types(
                 self.client.hexists(scanned_key, resource_name)
@@ -135,35 +112,35 @@ async def find_keys(self, resource: tuple[str, str]) -> list[UserSession]:
         field, value = resource
         return [
             self._decode_hash_key(hash_key)
-            async for hash_key in self.client.scan_iter(match=f"*:{_RESOURCE_SUFFIX}")
+            async for hash_key in self.client.scan_iter(match=f"*:{RESOURCE_SUFFIX}")
             if value
             == await handle_redis_returns_union_types(self.client.hget(hash_key, field))
         ]
 
     async def set_key_alive(self, key: UserSession, timeout: int) -> None:
         # setting the timeout to always expire, timeout > 0
         timeout = int(max(1, timeout))
-        hash_key = f"{self._hash_key(key)}:{_ALIVE_SUFFIX}"
+        hash_key = f"{key.to_redis_hash_key()}:{ALIVE_SUFFIX}"
         await self.client.set(hash_key, 1, ex=timeout)
 
     async def is_key_alive(self, key: UserSession) -> bool:
-        hash_key = f"{self._hash_key(key)}:{_ALIVE_SUFFIX}"
+        hash_key = f"{key.to_redis_hash_key()}:{ALIVE_SUFFIX}"
         return bool(await self.client.exists(hash_key) > 0)
 
     async def remove_key(self, key: UserSession) -> None:
         await self.client.delete(
-            f"{self._hash_key(key)}:{_RESOURCE_SUFFIX}",
-            f"{self._hash_key(key)}:{_ALIVE_SUFFIX}",
+            f"{key.to_redis_hash_key()}:{RESOURCE_SUFFIX}",
+            f"{key.to_redis_hash_key()}:{ALIVE_SUFFIX}",
         )
 
     async def get_all_resource_keys(self) -> tuple[AliveSessions, DeadSessions]:
         alive_keys = [
             self._decode_hash_key(hash_key)
-            async for hash_key in self.client.scan_iter(match=f"*:{_ALIVE_SUFFIX}")
+            async for hash_key in self.client.scan_iter(match=f"*:{ALIVE_SUFFIX}")
         ]
         dead_keys = [
             self._decode_hash_key(hash_key)
-            async for hash_key in self.client.scan_iter(match=f"*:{_RESOURCE_SUFFIX}")
+            async for hash_key in self.client.scan_iter(match=f"*:{RESOURCE_SUFFIX}")
             if self._decode_hash_key(hash_key) not in alive_keys
         ]
 

services/web/server/src/simcore_service_webserver/resource_manager/user_sessions.py

@@ -8,10 +8,9 @@
 from models_library.users import UserID
 from servicelib.logging_utils import get_log_record_extra, log_context
 
+from .models import ResourcesDict, UserSession
 from .registry import (
     RedisResourceRegistry,
-    ResourcesDict,
-    UserSession,
     get_registry,
 )
 from .settings import ResourceManagerSettings, get_plugin_settings
@@ -31,12 +30,6 @@ def _get_service_deletion_timeout(app: web.Application) -> int:
     return settings.RESOURCE_MANAGER_RESOURCE_TTL_S
 
 
-@dataclass(order=True, frozen=True)
-class UserSessionID:
-    user_id: UserID
-    client_session_id: str
-
-
 @dataclass
 class UserSessionResourcesRegistry:
     """
@@ -61,7 +54,7 @@ class UserSessionResourcesRegistry:
 
     """
 
-    user_id: int
+    user_id: UserID
     client_session_id: str | None  # Every tab that a user opens
     app: web.Application
 
@@ -71,7 +64,7 @@ def _registry(self) -> RedisResourceRegistry:
 
     def _resource_key(self) -> UserSession:
         return UserSession(
-            user_id=f"{self.user_id}",
+            user_id=self.user_id,
             client_session_id=self.client_session_id or "*",
         )
 
@@ -138,11 +131,10 @@ async def find_socket_ids(self) -> list[str]:
             extra=get_log_record_extra(user_id=self.user_id),
         )
 
-        user_sockets: list[str] = await self._registry.find_resources(
-            {"user_id": f"{self.user_id}", "client_session_id": "*"},
+        return await self._registry.find_resources(
+            UserSession(user_id=self.user_id, client_session_id="*"),
             _SOCKET_ID_FIELDNAME,
         )
-        return user_sockets
 
     async def find_all_resources_of_user(self, key: str) -> list[str]:
         with log_context(
@@ -151,10 +143,9 @@ async def find_all_resources_of_user(self, key: str) -> list[str]:
             msg=f"{self.user_id=} finding all {key} from registry",
             extra=get_log_record_extra(user_id=self.user_id),
         ):
-            resources: list[str] = await get_registry(self.app).find_resources(
-                {"user_id": f"{self.user_id}", "client_session_id": "*"}, key
+            return await get_registry(self.app).find_resources(
+                UserSession(user_id=self.user_id, client_session_id="*"), key
             )
-            return resources
 
     async def find(self, resource_name: str) -> list[str]:
         _logger.debug(
@@ -165,10 +156,7 @@ async def find(self, resource_name: str) -> list[str]:
             extra=get_log_record_extra(user_id=self.user_id),
         )
 
-        resource_values: list[str] = await self._registry.find_resources(
-            self._resource_key(), resource_name
-        )
-        return resource_values
+        return await self._registry.find_resources(self._resource_key(), resource_name)
 
     async def add(self, key: str, value: str) -> None:
         _logger.debug(
@@ -196,25 +184,15 @@ async def remove(self, key: str) -> None:
     @staticmethod
     async def find_users_of_resource(
         app: web.Application, key: str, value: str
-    ) -> list[UserSessionID]:
+    ) -> list[UserSession]:
         registry = get_registry(app)
-        registry_keys: list[UserSession] = await registry.find_keys(
-            resource=(key, value)
-        )
-        users_sessions_ids: list[UserSessionID] = [
-            UserSessionID(
-                user_id=int(r["user_id"]),
-                client_session_id=r["client_session_id"],
-            )
-            for r in registry_keys
-        ]
-        return users_sessions_ids
+        return await registry.find_keys(resource=(key, value))
 
-    def get_id(self) -> UserSessionID:
+    def get_id(self) -> UserSession:
         if self.client_session_id is None:
             msg = f"Cannot build UserSessionID with missing {self.client_session_id=}"
             raise ValueError(msg)
-        return UserSessionID(
+        return UserSession(
             user_id=self.user_id, client_session_id=self.client_session_id
         )