refactor(cli): implement filelock to protect db_path.

Author: Zhe Yu

pyproject.toml

@@ -20,6 +20,7 @@ dependencies = [
     "colorlog",
     "charset-normalizer>=3.4.1",
     "json5",
+    "filelock>=3.15.0",
 ]
 requires-python = ">=3.11,<3.14"
 readme = "README.md"

src/vectorcode/cli_utils.py

@@ -12,6 +12,7 @@
 
 import json5
 import shtab
+from filelock import AsyncFileLock
 
 from vectorcode import __version__
 
@@ -610,3 +611,18 @@ def config_logging(
         handlers=handlers,
         level=level,
     )
+
+
+class LockManager:
+    """
+    A class that manages file locks that protects the database files in daemon processes (LSP, MCP).
+    """
+
+    def __init__(self) -> None:
+        self.__locks: dict[str, AsyncFileLock] = {}
+
+    def get(self, path: str | os.PathLike) -> AsyncFileLock:
+        path = str(expand_path(str(path), True))
+        if self.__locks.get(path) is None:
+            self.__locks[path] = AsyncFileLock(path)  # pyright: ignore[reportArgumentType]
+        return self.__locks[path]

src/vectorcode/common.py

@@ -5,6 +5,7 @@
 import socket
 import subprocess
 import sys
+from asyncio.subprocess import Process
 from typing import Any, AsyncGenerator
 from urllib.parse import urlparse
 
@@ -261,3 +262,30 @@ async def list_collection_files(collection: AsyncCollection) -> list[str]:
             or []
         )
     )
+
+
+class ClientManager:
+    __singleton = None
+
+    # keys: project roots
+    # values: clients
+    __clients: dict[str, AsyncClientAPI] = {}
+    __server_processes = []
+
+    @classmethod
+    def get_instance(cls) -> "ClientManager":
+        if cls.__singleton is None:
+            cls.__singleton = ClientManager()
+        return cls.__singleton
+
+    async def get(self, configs: Config) -> AsyncClientAPI:
+        project_root = str(expand_path(str(configs.project_root), True))
+        if self.__clients.get(project_root) is None:
+            if not await try_server(configs.db_url):
+                self.__server_processes.append(await start_server(configs))
+
+            self.__clients[project_root] = await get_client(configs)
+        return self.__clients[project_root]
+
+    def get_processes(self) -> list[Process]:
+        return self.__server_processes

src/vectorcode/main.py

@@ -12,6 +12,7 @@
     get_project_config,
     parse_cli_args,
 )
+from vectorcode.common import ClientManager
 
 logger = logging.getLogger(name=__name__)
 
@@ -63,12 +64,6 @@ async def async_main():
 
             return await chunks(final_configs)
 
-    from vectorcode.common import start_server, try_server
-
-    server_process = None
-    if not await try_server(final_configs.db_url):
-        server_process = await start_server(final_configs)
-
     if final_configs.pipe:  # pragma: nocover
         # NOTE: NNCF (intel GPU acceleration for sentence transformer) keeps showing logs.
         # This disables logs below ERROR so that it doesn't hurt the `pipe` output.
@@ -105,10 +100,9 @@ async def async_main():
         return_val = 1
         logger.error(traceback.format_exc())
     finally:
-        if server_process is not None:
-            logger.info("Shutting down the bundled Chromadb instance.")
-            server_process.terminate()
-            await server_process.wait()
+        for p in ClientManager.get_instance().get_processes():
+            p.terminate()
+            await p.wait()
         return return_val
 
 

src/vectorcode/subcommands/query/__init__.py

@@ -17,7 +17,7 @@
     expand_path,
 )
 from vectorcode.common import (
-    get_client,
+    ClientManager,
     get_collection,
     verify_ef,
 )
@@ -160,7 +160,7 @@ async def query(configs: Config) -> int:
             "Having both chunk and document in the output is not supported!",
         )
         return 1
-    client = await get_client(configs)
+    client = await ClientManager.get_instance().get(configs)
     try:
         collection = await get_collection(client, configs, False)
         if not verify_ef(collection, configs):