refactor(cli): Use a context manager for client with filelock when necessary

Author: Davidyz

src/vectorcode/cli_utils.py

@@ -627,8 +627,14 @@ def __new__(cls) -> "LockManager":
             cls.__singleton.__locks = {}
         return cls.__singleton
 
-    def get(self, path: str | os.PathLike) -> AsyncFileLock:
+    def get_lock(self, path: str | os.PathLike) -> AsyncFileLock:
         path = str(expand_path(str(path), True))
+        if os.path.isdir(path):
+            lock_file = os.path.join(path, "vectorcode.lock")
+            logger.info(f"Creating {lock_file} for locking.")
+            with open(lock_file, mode="w") as fin:
+                fin.write("")
+            path = lock_file
         if self.__locks.get(path) is None:
             self.__locks[path] = AsyncFileLock(path)  # pyright: ignore[reportArgumentType]
         return self.__locks[path]

src/vectorcode/common.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 import hashlib
 import logging
 import os
@@ -18,7 +19,7 @@
 from chromadb.config import APIVersion, Settings
 from chromadb.utils import embedding_functions
 
-from vectorcode.cli_utils import Config, expand_path
+from vectorcode.cli_utils import Config, LockManager, expand_path
 
 logger = logging.getLogger(name=__name__)
 
@@ -114,26 +115,6 @@ async def start_server(configs: Config):
     return process
 
 
-async def get_client(configs: Config) -> AsyncClientAPI:
-    settings: dict[str, Any] = {"anonymized_telemetry": False}
-    if isinstance(configs.db_settings, dict):
-        valid_settings = {
-            k: v for k, v in configs.db_settings.items() if k in Settings.__fields__
-        }
-        settings.update(valid_settings)
-    parsed_url = urlparse(configs.db_url)
-    settings["chroma_server_host"] = parsed_url.hostname or "127.0.0.1"
-    settings["chroma_server_http_port"] = parsed_url.port or 8000
-    settings["chroma_server_ssl_enabled"] = parsed_url.scheme == "https"
-    settings["chroma_server_api_default_path"] = parsed_url.path or APIVersion.V2
-    settings_obj = Settings(**settings)
-    return await chromadb.AsyncHttpClient(
-        settings=settings_obj,
-        host=str(settings_obj.chroma_server_host),
-        port=int(settings_obj.chroma_server_http_port or 8000),
-    )
-
-
 def get_collection_name(full_path: str) -> str:
     full_path = str(expand_path(full_path, absolute=True))
     hasher = hashlib.sha256()
@@ -276,27 +257,61 @@ def __new__(cls) -> "ClientManager":
             cls.__singleton.__clients = {}
         return cls.__singleton
 
-    async def get_client(self, configs: Config) -> _ClientModel:
+    @contextlib.asynccontextmanager
+    async def get_client(self, configs: Config, need_lock: bool = True):
         project_root = str(expand_path(str(configs.project_root), True))
+        is_bundled = False
         if self.__clients.get(project_root) is None:
-            is_bundled = False
             process = None
             if not await try_server(configs.db_url):
                 logger.info(f"Starting a new server at {configs.db_url}")
                 process = await start_server(configs)
                 is_bundled = True
 
             self.__clients[project_root] = _ClientModel(
-                client=await get_client(configs), is_bundled=is_bundled, process=process
+                client=await self._create_client(configs),
+                is_bundled=is_bundled,
+                process=process,
             )
-        return self.__clients[project_root]
+        lock = None
+        if self.__clients[project_root].is_bundled and need_lock:
+            lock = LockManager().get_lock(str(configs.db_path))
+            logger.debug(f"Locking {configs.db_path}")
+            await lock.acquire()
+        yield self.__clients[project_root].client
+        if lock is not None:
+            logger.debug(f"Unlocking {configs.db_path}")
+            await lock.release()
 
     def get_processes(self) -> list[Process]:
         return [i.process for i in self.__clients.values() if i.process is not None]
 
     async def kill_servers(self):
         termination_tasks: list[asyncio.Task] = []
         for p in ClientManager().get_processes():
+            logger.info(f"Killing bundled chroma server with PID: {p.pid}")
             p.terminate()
             termination_tasks.append(asyncio.create_task(p.wait()))
         await asyncio.gather(*termination_tasks)
+
+    async def _create_client(self, configs: Config) -> AsyncClientAPI:
+        settings: dict[str, Any] = {"anonymized_telemetry": False}
+        if isinstance(configs.db_settings, dict):
+            valid_settings = {
+                k: v for k, v in configs.db_settings.items() if k in Settings.__fields__
+            }
+            settings.update(valid_settings)
+        parsed_url = urlparse(configs.db_url)
+        settings["chroma_server_host"] = parsed_url.hostname or "127.0.0.1"
+        settings["chroma_server_http_port"] = parsed_url.port or 8000
+        settings["chroma_server_ssl_enabled"] = parsed_url.scheme == "https"
+        settings["chroma_server_api_default_path"] = parsed_url.path or APIVersion.V2
+        settings_obj = Settings(**settings)
+        return await chromadb.AsyncHttpClient(
+            settings=settings_obj,
+            host=str(settings_obj.chroma_server_host),
+            port=int(settings_obj.chroma_server_http_port or 8000),
+        )
+
+    def clear(self):
+        self.__clients.clear()