Changed Lock primites to asyncio

Author: aponcedeleonch

src/codegate/muxing/router.py

@@ -24,7 +24,6 @@ def __init__(self, provider_registry: ProviderRegistry):
         self._setup_routes()
         self._provider_registry = provider_registry
         self._response_adapter = ResponseAdapter()
-        self._mux_registry = rulematcher.get_muxing_rules_registry()
 
     @property
     def route_name(self) -> str:
@@ -55,10 +54,11 @@ async def route_to_dest_provider(
             body = await request.body()
             data = json.loads(body)
 
+            mux_registry = await rulematcher.get_muxing_rules_registry()
             try:
                 # TODO: For future releases we will have to idenify a thing_to_match
                 # and use our registry to get the correct muxes for the active workspace
-                model_route = self._mux_registry.get_match_for_active_workspace(thing_to_match=None)
+                model_route = await mux_registry.get_match_for_active_workspace(thing_to_match=None)
             except Exception as e:
                 logger.error(f"Error getting active workspace muxes: {e}")
                 raise HTTPException(str(e))

src/codegate/muxing/rulematcher.py

@@ -1,7 +1,6 @@
 import copy
 from abc import ABC, abstractmethod
-from collections import UserDict
-from threading import Lock, RLock
+from asyncio import Lock
 from typing import List, Optional
 
 from codegate.db import models as db_models
@@ -11,13 +10,13 @@
 _singleton_lock = Lock()
 
 
-def get_muxing_rules_registry():
+async def get_muxing_rules_registry():
     """Returns a singleton instance of the muxing rules registry."""
 
     global _muxrules_sgtn
 
     if _muxrules_sgtn is None:
-        with _singleton_lock:
+        async with _singleton_lock:
             if _muxrules_sgtn is None:
                 _muxrules_sgtn = MuxingRulesinWorkspaces()
 
@@ -79,38 +78,47 @@ def match(self, thing_to_match) -> bool:
         return True
 
 
-class MuxingRulesinWorkspaces(UserDict):
+class MuxingRulesinWorkspaces:
     """A thread safe dictionary to store the muxing rules in workspaces."""
 
     def __init__(self) -> None:
         super().__init__()
-        self._lock = RLock()
+        self._lock = Lock()
         self._active_workspace = ""
+        self._ws_rules = {}
 
-    def __getitem__(self, key: str) -> List[MuxingRuleMatcher]:
-        with self._lock:
-            # We return a copy so concurrent modifications don't affect the original
-            return copy.deepcopy(super().__getitem__(key))
+    async def get_ws_rules(self, workspace_name: str) -> List[MuxingRuleMatcher]:
+        """Get the rules for the given workspace."""
+        async with self._lock:
+            return copy.deepcopy(self._ws_rules.get(workspace_name, []))
 
-    def __setitem__(self, key: str, value: List[MuxingRuleMatcher]) -> None:
-        with self._lock:
-            super().__setitem__(key, value)
+    async def set_ws_rules(self, workspace_name: str, rules: List[MuxingRuleMatcher]) -> None:
+        """Set the rules for the given workspace."""
+        async with self._lock:
+            self._ws_rules[workspace_name] = rules
 
-    def __delitem__(self, key: str) -> None:
-        with self._lock:
-            super().__delitem__(key)
+    async def delete_ws_rules(self, workspace_name: str) -> None:
+        """Delete the rules for the given workspace."""
+        async with self._lock:
+            del self._ws_rules[workspace_name]
 
-    def set_active_workspace(self, workspace_name: str) -> None:
+    async def set_active_workspace(self, workspace_name: str) -> None:
         """Set the active workspace."""
         self._active_workspace = workspace_name
 
-    def get_match_for_active_workspace(self, thing_to_match) -> Optional[ModelRoute]:
+    async def get_registries(self) -> List[str]:
+        """Get the list of workspaces."""
+        async with self._lock:
+            return list(self._ws_rules.keys())
+
+    async def get_match_for_active_workspace(self, thing_to_match) -> Optional[ModelRoute]:
         """Get the first match for the given thing_to_match."""
 
         # We iterate over all the rules and return the first match
         # Since we already do a deepcopy in __getitem__, we don't need to lock here
         try:
-            for rule in self[self._active_workspace]:
+            rules = await self.get_ws_rules(self._active_workspace)
+            for rule in rules:
                 if rule.match(thing_to_match):
                     return rule.destination()
             return None

src/codegate/workspaces/crud.py

@@ -37,12 +37,8 @@ class WorkspaceMuxRuleDoesNotExistError(WorkspaceCrudError):
 
 class WorkspaceCrud:
 
-    def __init__(
-        self,
-        mux_registry: rulematcher.MuxingRulesinWorkspaces = rulematcher.get_muxing_rules_registry(),
-    ):
+    def __init__(self):
         self._db_reader = DbReader()
-        self._mux_registry = mux_registry
 
     async def add_workspace(self, new_workspace_name: str) -> WorkspaceRow:
         """
@@ -141,7 +137,8 @@ async def activate_workspace(self, workspace_name: str):
         await db_recorder.update_session(session)
 
         # Ensure the mux registry is updated
-        self._mux_registry.set_active_workspace(workspace.name)
+        mux_registry = await rulematcher.get_muxing_rules_registry()
+        await mux_registry.set_active_workspace(workspace.name)
         return
 
     async def recover_workspace(self, workspace_name: str):
@@ -198,7 +195,8 @@ async def soft_delete_workspace(self, workspace_name: str):
             raise WorkspaceCrudError(f"Error deleting workspace {workspace_name}")
 
         # Remove the muxes from the registry
-        del self._mux_registry[workspace_name]
+        mux_registry = await rulematcher.get_muxing_rules_registry()
+        await mux_registry.delete_ws_rules(workspace_name)
         return
 
     async def hard_delete_workspace(self, workspace_name: str):
@@ -293,7 +291,8 @@ async def set_muxes(self, workspace_name: str, muxes):
             priority += 1
 
         # Set routing list for the workspace
-        self._mux_registry[workspace_name] = matchers
+        mux_registry = await rulematcher.get_muxing_rules_registry()
+        await mux_registry.set_ws_rules(workspace_name, matchers)
 
     async def get_routing_for_mux(self, mux) -> rulematcher.ModelRoute:
         """Get the routing for a mux
@@ -359,7 +358,8 @@ async def initialize_mux_registry(self) -> None:
 
         active_ws = await self.get_active_workspace()
         if active_ws:
-            self._mux_registry.set_active_workspace(active_ws.name)
+            mux_registry = await rulematcher.get_muxing_rules_registry()
+            await mux_registry.set_active_workspace(active_ws.name)
 
         await self.repopulate_mux_cache()
 
@@ -369,12 +369,14 @@ async def repopulate_mux_cache(self) -> None:
         # Get all workspaces
         workspaces = await self.get_workspaces()
 
+        mux_registry = await rulematcher.get_muxing_rules_registry()
+
         # Remove any workspaces from cache that are not in the database
         ws_names = set(ws.name for ws in workspaces)
-        cached_ws = set(self._mux_registry.keys())
+        cached_ws = set(await mux_registry.get_registries())
         ws_to_remove = cached_ws - ws_names
         for ws in ws_to_remove:
-            del self._mux_registry[ws]
+            await mux_registry.delete_ws_rules(ws)
 
         # For each workspace, get the muxes and set them in the registry
         for ws in workspaces:
@@ -386,4 +388,4 @@ async def repopulate_mux_cache(self) -> None:
                 route = await self.get_routing_for_db_mux(mux)
                 matchers.append(rulematcher.MuxingMatcherFactory.create(mux, route))
 
-            self._mux_registry[ws.name] = matchers
+            await mux_registry.set_ws_rules(ws.name, matchers)