Skip to content
This repository was archived by the owner on Jun 5, 2025. It is now read-only.

Commit 1f11171

Browse files
aponcedeleonchaponcedeleonch
committed
Changed Lock primites to asyncio
1 parent 0bf51ed commit 1f11171

3 files changed

Lines changed: 43 additions & 33 deletions

File tree

‎src/codegate/muxing/router.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ def __init__(self, provider_registry: ProviderRegistry):
2424
self._setup_routes()
2525
self._provider_registry = provider_registry
2626
self._response_adapter = ResponseAdapter()
27-
self._mux_registry = rulematcher.get_muxing_rules_registry()
2827

2928
@property
3029
def route_name(self) -> str:
@@ -55,10 +54,11 @@ async def route_to_dest_provider(
5554
body = await request.body()
5655
data = json.loads(body)
5756

57+
mux_registry = await rulematcher.get_muxing_rules_registry()
5858
try:
5959
# TODO: For future releases we will have to idenify a thing_to_match
6060
# and use our registry to get the correct muxes for the active workspace
61-
model_route = self._mux_registry.get_match_for_active_workspace(thing_to_match=None)
61+
model_route = await mux_registry.get_match_for_active_workspace(thing_to_match=None)
6262
except Exception as e:
6363
logger.error(f"Error getting active workspace muxes: {e}")
6464
raise HTTPException(str(e))

‎src/codegate/muxing/rulematcher.py‎

Lines changed: 27 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import copy
22
from abc import ABC, abstractmethod
3-
from collections import UserDict
4-
from threading import Lock, RLock
3+
from asyncio import Lock
54
from typing import List, Optional
65

76
from codegate.db import models as db_models
@@ -11,13 +10,13 @@
1110
_singleton_lock = Lock()
1211

1312

14-
def get_muxing_rules_registry():
13+
async def get_muxing_rules_registry():
1514
"""Returns a singleton instance of the muxing rules registry."""
1615

1716
global _muxrules_sgtn
1817

1918
if _muxrules_sgtn is None:
20-
with _singleton_lock:
19+
async with _singleton_lock:
2120
if _muxrules_sgtn is None:
2221
_muxrules_sgtn = MuxingRulesinWorkspaces()
2322

@@ -79,38 +78,47 @@ def match(self, thing_to_match) -> bool:
7978
return True
8079

8180

82-
class MuxingRulesinWorkspaces(UserDict):
81+
class MuxingRulesinWorkspaces:
8382
"""A thread safe dictionary to store the muxing rules in workspaces."""
8483

8584
def __init__(self) -> None:
8685
super().__init__()
87-
self._lock = RLock()
86+
self._lock = Lock()
8887
self._active_workspace = ""
88+
self._ws_rules = {}
8989

90-
def __getitem__(self, key: str) -> List[MuxingRuleMatcher]:
91-
with self._lock:
92-
# We return a copy so concurrent modifications don't affect the original
93-
return copy.deepcopy(super().__getitem__(key))
90+
async def get_ws_rules(self, workspace_name: str) -> List[MuxingRuleMatcher]:
91+
"""Get the rules for the given workspace."""
92+
async with self._lock:
93+
return copy.deepcopy(self._ws_rules.get(workspace_name, []))
9494

95-
def __setitem__(self, key: str, value: List[MuxingRuleMatcher]) -> None:
96-
with self._lock:
97-
super().__setitem__(key, value)
95+
async def set_ws_rules(self, workspace_name: str, rules: List[MuxingRuleMatcher]) -> None:
96+
"""Set the rules for the given workspace."""
97+
async with self._lock:
98+
self._ws_rules[workspace_name] = rules
9899

99-
def __delitem__(self, key: str) -> None:
100-
with self._lock:
101-
super().__delitem__(key)
100+
async def delete_ws_rules(self, workspace_name: str) -> None:
101+
"""Delete the rules for the given workspace."""
102+
async with self._lock:
103+
del self._ws_rules[workspace_name]
102104

103-
def set_active_workspace(self, workspace_name: str) -> None:
105+
async def set_active_workspace(self, workspace_name: str) -> None:
104106
"""Set the active workspace."""
105107
self._active_workspace = workspace_name
106108

107-
def get_match_for_active_workspace(self, thing_to_match) -> Optional[ModelRoute]:
109+
async def get_registries(self) -> List[str]:
110+
"""Get the list of workspaces."""
111+
async with self._lock:
112+
return list(self._ws_rules.keys())
113+
114+
async def get_match_for_active_workspace(self, thing_to_match) -> Optional[ModelRoute]:
108115
"""Get the first match for the given thing_to_match."""
109116

110117
# We iterate over all the rules and return the first match
111118
# Since we already do a deepcopy in __getitem__, we don't need to lock here
112119
try:
113-
for rule in self[self._active_workspace]:
120+
rules = await self.get_ws_rules(self._active_workspace)
121+
for rule in rules:
114122
if rule.match(thing_to_match):
115123
return rule.destination()
116124
return None

‎src/codegate/workspaces/crud.py‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -37,12 +37,8 @@ class WorkspaceMuxRuleDoesNotExistError(WorkspaceCrudError):
3737

3838
class WorkspaceCrud:
3939

40-
def __init__(
41-
self,
42-
mux_registry: rulematcher.MuxingRulesinWorkspaces = rulematcher.get_muxing_rules_registry(),
43-
):
40+
def __init__(self):
4441
self._db_reader = DbReader()
45-
self._mux_registry = mux_registry
4642

4743
async def add_workspace(self, new_workspace_name: str) -> WorkspaceRow:
4844
"""
@@ -141,7 +137,8 @@ async def activate_workspace(self, workspace_name: str):
141137
await db_recorder.update_session(session)
142138

143139
# Ensure the mux registry is updated
144-
self._mux_registry.set_active_workspace(workspace.name)
140+
mux_registry = await rulematcher.get_muxing_rules_registry()
141+
await mux_registry.set_active_workspace(workspace.name)
145142
return
146143

147144
async def recover_workspace(self, workspace_name: str):
@@ -198,7 +195,8 @@ async def soft_delete_workspace(self, workspace_name: str):
198195
raise WorkspaceCrudError(f"Error deleting workspace {workspace_name}")
199196

200197
# Remove the muxes from the registry
201-
del self._mux_registry[workspace_name]
198+
mux_registry = await rulematcher.get_muxing_rules_registry()
199+
await mux_registry.delete_ws_rules(workspace_name)
202200
return
203201

204202
async def hard_delete_workspace(self, workspace_name: str):
@@ -293,7 +291,8 @@ async def set_muxes(self, workspace_name: str, muxes):
293291
priority += 1
294292

295293
# Set routing list for the workspace
296-
self._mux_registry[workspace_name] = matchers
294+
mux_registry = await rulematcher.get_muxing_rules_registry()
295+
await mux_registry.set_ws_rules(workspace_name, matchers)
297296

298297
async def get_routing_for_mux(self, mux) -> rulematcher.ModelRoute:
299298
"""Get the routing for a mux
@@ -359,7 +358,8 @@ async def initialize_mux_registry(self) -> None:
359358

360359
active_ws = await self.get_active_workspace()
361360
if active_ws:
362-
self._mux_registry.set_active_workspace(active_ws.name)
361+
mux_registry = await rulematcher.get_muxing_rules_registry()
362+
await mux_registry.set_active_workspace(active_ws.name)
363363

364364
await self.repopulate_mux_cache()
365365

@@ -369,12 +369,14 @@ async def repopulate_mux_cache(self) -> None:
369369
# Get all workspaces
370370
workspaces = await self.get_workspaces()
371371

372+
mux_registry = await rulematcher.get_muxing_rules_registry()
373+
372374
# Remove any workspaces from cache that are not in the database
373375
ws_names = set(ws.name for ws in workspaces)
374-
cached_ws = set(self._mux_registry.keys())
376+
cached_ws = set(await mux_registry.get_registries())
375377
ws_to_remove = cached_ws - ws_names
376378
for ws in ws_to_remove:
377-
del self._mux_registry[ws]
379+
await mux_registry.delete_ws_rules(ws)
378380

379381
# For each workspace, get the muxes and set them in the registry
380382
for ws in workspaces:
@@ -386,4 +388,4 @@ async def repopulate_mux_cache(self) -> None:
386388
route = await self.get_routing_for_db_mux(mux)
387389
matchers.append(rulematcher.MuxingMatcherFactory.create(mux, route))
388390

389-
self._mux_registry[ws.name] = matchers
391+
await mux_registry.set_ws_rules(ws.name, matchers)

0 commit comments

Comments
 (0)