diff --git a/src/cocoindex_code/client.py b/src/cocoindex_code/client.py index f65dafb..f93924d 100644 --- a/src/cocoindex_code/client.py +++ b/src/cocoindex_code/client.py @@ -13,7 +13,8 @@ import subprocess import sys import time -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from multiprocessing.connection import Client, Connection from pathlib import Path from typing import NamedTuple @@ -125,6 +126,49 @@ def _is_daemon_supervised() -> bool: return os.environ.get("COCOINDEX_CODE_DAEMON_SUPERVISED") == "1" +@contextmanager +def _daemon_start_lock() -> Iterator[None]: + """Serialize daemon startup across independent client processes. + + Each MCP server is a separate process, so the process-local + ``_daemon_ensured`` flag cannot prevent a startup stampede. Keep the lock + file itself after releasing the OS lock: unlinking it would let a new + opener lock a different inode while another process still holds the old + one. + """ + daemon_runtime_dir().mkdir(parents=True, exist_ok=True) + lock_path = daemon_runtime_dir() / "daemon-start.lock" + + with open(lock_path, "a+b") as lock_file: + if sys.platform == "win32": + import msvcrt + + lock_file.seek(0) + if lock_file.read(1) == b"": + lock_file.write(b"\0") + lock_file.flush() + lock_file.seek(0) + while True: + try: + msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1) + break + except OSError: + time.sleep(0.1) + else: + import fcntl + + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) + + try: + yield + finally: + if sys.platform == "win32": + lock_file.seek(0) + msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1) + else: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) + + def _connect_and_handshake() -> Connection: """Connect to the daemon and perform the version handshake. @@ -142,6 +186,7 @@ def _connect_and_handshake() -> Connection: """ global _daemon_ensured, _ensured_daemon_pid, _consecutive_crash_restarts # noqa: PLW0603 + startup_error: Exception try: result = _raw_connect_and_handshake() _daemon_ensured = True @@ -149,44 +194,77 @@ def _connect_and_handshake() -> Connection: # The daemon answered without needing a restart — any crash streak is over. _consecutive_crash_restarts = 0 return result.conn - except DaemonVersionError as e: + except DaemonVersionError as error: + startup_error = error # `resp.ok` is False only for a real version mismatch. Once we have # ensured a matching daemon, a fresh version mismatch means the binary - # was swapped under us — fail fast. A settings-only restart request - # (resp.ok True, but the loaded settings mtime moved) is expected; - # restart the daemon below so it reloads them. (stop_daemon sends a - # StopRequest, so this restart leaves a graceful-exit marker and the - # next connect stays silent.) - if _daemon_ensured and not e.resp.ok: + # was swapped under us — fail fast. + if _daemon_ensured and not error.resp.ok: raise - stop_daemon() - except DaemonProtocolError: + except DaemonProtocolError as error: + startup_error = error # The reply didn't even decode — a daemon whose wire protocol # drifted beyond what the version-mismatch path can express (e.g. - # still running a pre-upgrade binary). Same treatment as a version - # mismatch: restart it below; once a matching daemon was already - # ensured this cannot legitimately happen, so fail fast instead of - # looping on restarts. + # still running a pre-upgrade binary). if _daemon_ensured: raise - stop_daemon() - except (ConnectionRefusedError, OSError): - # No daemon answered. Normal on the first call of this process (start - # one below). If we had already ensured one, decide crash vs graceful - # exit — a crash restarts loudly and is capped, a graceful exit (e.g. - # a manual stop) restarts silently. - if _daemon_ensured: - _handle_vanished_daemon() + except (ConnectionRefusedError, OSError) as error: + startup_error = error if _is_daemon_supervised(): # Supervisor is responsible for (re)starting the daemon — just wait # for the socket to reappear. + _prepare_daemon_restart(startup_error) _wait_for_daemon() - else: + return _connect_after_daemon_start() + + # Multiple MCP servers can arrive here together. The first process starts + # the daemon while holding the OS lock. Every waiter rechecks after + # acquiring it and joins that daemon instead of starting another one. + with _daemon_start_lock(): + try: + result = _raw_connect_and_handshake() + except ( + DaemonVersionError, + DaemonProtocolError, + ConnectionRefusedError, + OSError, + ) as locked_error: + _prepare_daemon_restart(locked_error) + else: + _daemon_ensured = True + _ensured_daemon_pid = result.resp.pid + _consecutive_crash_restarts = 0 + return result.conn + proc = start_daemon() _wait_for_daemon(proc=proc) + return _connect_after_daemon_start() + + +def _prepare_daemon_restart(error: Exception) -> None: + """Apply restart policy after a failed handshake has been confirmed.""" + if isinstance(error, DaemonVersionError): + if _daemon_ensured and not error.resp.ok: + raise error + # A settings-only mismatch is expected. Stop the daemon so the new + # process reloads global_settings.yml. + stop_daemon() + elif isinstance(error, DaemonProtocolError): + if _daemon_ensured: + raise error + stop_daemon() + elif _daemon_ensured: + # No daemon answered. Decide crash vs graceful exit only after the + # startup lock recheck, so a transient backlog refusal is not reported + # as a crash. + _handle_vanished_daemon() + + +def _connect_after_daemon_start() -> Connection: + """Wait for a newly started daemon to complete a real handshake.""" + global _daemon_ensured, _ensured_daemon_pid # noqa: PLW0603 - # Verify the fresh daemon is reachable for _attempt in range(10): try: result = _raw_connect_and_handshake() diff --git a/src/cocoindex_code/daemon.py b/src/cocoindex_code/daemon.py index d2b38ab..e1a0b79 100644 --- a/src/cocoindex_code/daemon.py +++ b/src/cocoindex_code/daemon.py @@ -693,7 +693,10 @@ def run_daemon( except Exception: pass - listener = Listener(sock_path, family=connection_family()) + # multiprocessing.Listener defaults to backlog=1. Several MCP clients can + # connect at once when an agent session starts, so a larger queue prevents + # transient refusals from being mistaken for a vanished daemon. + listener = Listener(sock_path, family=connection_family(), backlog=64) logger.info("Listening on %s", sock_path) loop = asyncio.new_event_loop() diff --git a/tests/test_client.py b/tests/test_client.py index 5b1f453..ff7090b 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -2,10 +2,13 @@ from __future__ import annotations +import multiprocessing +import os +import subprocess import tempfile from multiprocessing.connection import Connection from pathlib import Path -from typing import cast +from typing import Any, cast import pytest @@ -14,6 +17,86 @@ from cocoindex_code.protocol import HandshakeResponse +def _concurrent_connect_worker( + runtime_dir: str, + barrier: Any, + starts_path: str, + ready_path: str, + result_queue: Any, +) -> None: + """Run one isolated client process in the daemon-start stampede harness.""" + os.environ["COCOINDEX_CODE_RUNTIME_DIR"] = runtime_dir + + from cocoindex_code import client as worker_client + + raw_calls = 0 + sentinel_conn = cast(Connection, object()) + ok_resp = HandshakeResponse(ok=True, daemon_version="test", pid=42) + + def fake_raw() -> worker_client._HandshakeResult: + nonlocal raw_calls + raw_calls += 1 + if raw_calls == 1: + barrier.wait(timeout=10) + raise ConnectionRefusedError("simultaneous initial miss") + if not Path(ready_path).exists(): + raise ConnectionRefusedError("daemon is not ready") + return worker_client._HandshakeResult(conn=sentinel_conn, resp=ok_resp) + + def fake_start() -> subprocess.Popen[bytes]: + with Path(starts_path).open("a") as starts_file: + starts_file.write(f"{os.getpid()}\n") + Path(ready_path).touch() + return cast(subprocess.Popen[bytes], object()) + + worker_client._raw_connect_and_handshake = fake_raw + worker_client.start_daemon = fake_start + worker_client._wait_for_daemon = lambda **_kwargs: None + worker_client._is_daemon_supervised = lambda: False + worker_client._daemon_ensured = False + + try: + worker_client._connect_and_handshake() + except Exception as exc: + result_queue.put(f"{type(exc).__name__}: {exc}") + else: + result_queue.put(None) + + +def test_concurrent_clients_start_only_one_daemon(tmp_path: Path) -> None: + """Independent MCP processes must serialize and deduplicate daemon startup.""" + process_count = 6 + ctx = multiprocessing.get_context("spawn") + barrier = ctx.Barrier(process_count) + result_queue = ctx.Queue() + runtime_dir = tmp_path / "runtime" + runtime_dir.mkdir() + starts_path = tmp_path / "starts" + ready_path = tmp_path / "ready" + + processes = [ + ctx.Process( + target=_concurrent_connect_worker, + args=( + str(runtime_dir), + barrier, + str(starts_path), + str(ready_path), + result_queue, + ), + ) + for _ in range(process_count) + ] + for process in processes: + process.start() + for process in processes: + process.join(timeout=15) + + assert all(process.exitcode == 0 for process in processes) + assert [result_queue.get(timeout=1) for _ in processes] == [None] * process_count + assert len(starts_path.read_text().splitlines()) == 1 + + def test_client_connect_refuses_when_no_daemon( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -100,7 +183,7 @@ def test_connect_restarts_ensured_daemon_on_stale_settings( def fake_raw() -> client._HandshakeResult: calls["raw"] += 1 - if calls["raw"] == 1: + if calls["raw"] <= 2: raise client.DaemonVersionError( HandshakeResponse(ok=True, daemon_version="v1", pid=42, global_settings_mtime_us=1) ) @@ -117,7 +200,7 @@ def fake_raw() -> client._HandshakeResult: assert conn is sentinel_conn assert calls["stop"] == 1 # old daemon stopped assert calls["start"] == 1 # fresh daemon started to reload settings - assert calls["raw"] == 2 # reconnected after restart + assert calls["raw"] == 3 # lock recheck, then reconnect after restart def test_connect_fails_fast_on_version_mismatch_after_ensured( @@ -163,7 +246,7 @@ def test_connect_restarts_daemon_on_undecodable_handshake( def fake_raw() -> client._HandshakeResult: calls["raw"] += 1 - if calls["raw"] == 1: + if calls["raw"] <= 2: raise client.DaemonProtocolError("Undecodable handshake reply from daemon") return client._HandshakeResult(conn=cast(Connection, sentinel_conn), resp=ok_resp) @@ -178,7 +261,7 @@ def fake_raw() -> client._HandshakeResult: assert conn is sentinel_conn assert calls["stop"] == 1 # incompatible daemon stopped assert calls["start"] == 1 # fresh daemon started - assert calls["raw"] == 2 # reconnected after restart + assert calls["raw"] == 3 # lock recheck, then reconnect after restart def test_connect_fails_fast_on_undecodable_handshake_after_ensured( @@ -257,7 +340,7 @@ def _setup_vanished_daemon( def fake_raw() -> client._HandshakeResult: calls["raw"] += 1 - if calls["raw"] == 1: + if calls["raw"] <= 2: raise ConnectionRefusedError("daemon socket not found") return client._HandshakeResult(conn=cast(Connection, sentinel_conn), resp=ok_resp)