Coverage for src/lilbee/runtime/lock.py: 100%

1"""Cross-process locks: LanceDB write locking and the server singleton. 

2 

3Write locking combines an in-process mutex with a cross-process file lock 

4(filelock) so separate processes also coordinate writes. Read consistency is 

5handled by LanceDB's built-in MVCC via ``read_consistency_interval`` in 

6``lilbee.data.store``. The server lock makes ``lilbee serve`` a singleton per 

7data dir. 

8""" 

9 

10import json 

11import logging 

12import threading 

13import time 

14from collections.abc import Generator 

15from contextlib import contextmanager 

16from dataclasses import dataclass 

17from pathlib import Path 

18 

19from filelock import FileLock 

20from filelock import Timeout as FileLockTimeout 

21 

22from lilbee.core.config import cfg 

23 

24log = logging.getLogger(__name__) 

25 

26# Default timeout (seconds) for acquiring the write lock 

27LOCK_TIMEOUT = 30.0 

28# Grace (seconds) for a dying predecessor to release the server lock during a 

29# restart handoff before a new `lilbee serve` gives up. This budgets the PROMPT 

30# path only: a predecessor whose llama-swaps honor SIGTERM releases well inside 

31# it. It deliberately does NOT cover SIGKILL escalation -- the teardown that 

32# actually holds the lock (_release_engines -> stop_engine -> _stop_stale_swap in 

33# lilbee.providers.fleet.swap_manager) can spend _ORPHAN_STOP_TIMEOUT_S plus the 

34# kill and reap waits per group, i.e. tens of seconds across four groups on a 

35# wedged machine. Sizing this to that worst case would make every ordinary 

36# restart wait on a pathological one; instead the successor exits with 

37# LOCK_REFUSAL_EXIT_CODE and the operator retries. 

38SERVER_LOCK_TIMEOUT = 15.0 

39_SERVER_LOCK_NAME = "server.lock" 

40_SCOPE_LOCK_NAME = "server.scope.lock" 

41_SCOPE_OWNER_NAME = "server.scope.owner.json" 

42# Minimum blocking wait granted to the in-process mutex even when the file lock 

43# consumed the whole budget, so a deadline-edge acquire still gets a real attempt. 

44_MUTEX_MIN_WAIT = 0.1 

45 

46 

47class LockTimeoutError(TimeoutError): 

48 """Raised when a lock cannot be acquired within the timeout.""" 

49 

50 

51# In-process write mutex: serializes writers within the same process 

52_write_mutex = threading.Lock() 

53 

54 

55def _lock_path(lancedb_dir: Path | None) -> Path: 

56 return (lancedb_dir if lancedb_dir is not None else cfg.lancedb_dir) / ".lock" 

57 

58 

59def server_lock_path(data_dir: Path) -> Path: 

60 """Path of the one-server-per-data-dir lock file.""" 

61 return data_dir / _SERVER_LOCK_NAME 

62 

63 

64def acquire_server_lock(data_dir: Path, timeout: float = SERVER_LOCK_TIMEOUT) -> FileLock | None: 

65 """Hold the one-server-per-data-dir lock, or None when a live server owns it. 

66 

67 The lock is an OS file lock, so the kernel releases it the moment its holder 

68 exits, however it died; a crashed or killed server leaves no stale state. 

69 """ 

70 data_dir.mkdir(parents=True, exist_ok=True) 

71 lock = FileLock(server_lock_path(data_dir)) 

72 try: 

73 lock.acquire(timeout=timeout) 

74 except FileLockTimeout: 

75 return None 

76 return lock 

77 

78 

79@dataclass(frozen=True) 

80class ScopeOwner: 

81 """The data dir the server holding a scope lock is serving, for the refusal message.""" 

82 

83 data_dir: str 

84 

85 

86@dataclass(frozen=True) 

87class ScopeHold: 

88 """A held scope lock plus its owner sidecar; release removes both.""" 

89 

90 lock: FileLock 

91 owner_path: Path 

92 

93 def release(self) -> None: 

94 """Remove the owner sidecar, then free the scope for the next server.""" 

95 self.owner_path.unlink(missing_ok=True) 

96 self.lock.release() 

97 

98 

99def acquire_scope_lock( 

100 scope_dir: Path, data_dir: Path, timeout: float = SERVER_LOCK_TIMEOUT 

101) -> ScopeHold | None: 

102 """Hold the one-server-per-scope lock, or None when a live server owns the scope. 

103 

104 The scope is a directory shared by several would-be servers (the Obsidian 

105 plugin's shared root). Like the data-dir lock, the OS releases it the moment 

106 the holder exits. The owner sidecar records which data dir the holder is 

107 serving so a refused starter can name it in its message. 

108 """ 

109 scope_dir.mkdir(parents=True, exist_ok=True) 

110 lock = FileLock(scope_dir / _SCOPE_LOCK_NAME) 

111 try: 

112 lock.acquire(timeout=timeout) 

113 except FileLockTimeout: 

114 return None 

115 owner_path = scope_dir / _SCOPE_OWNER_NAME 

116 owner_path.write_text(json.dumps({"data_dir": str(data_dir)}), encoding="utf-8") 

117 return ScopeHold(lock, owner_path) 

118 

119 

120def read_scope_owner(scope_dir: Path) -> ScopeOwner | None: 

121 """The scope's recorded owner, or None when absent or unreadable.""" 

122 try: 

123 payload = json.loads((scope_dir / _SCOPE_OWNER_NAME).read_text(encoding="utf-8")) 

124 return ScopeOwner(data_dir=str(payload["data_dir"])) 

125 except (OSError, ValueError, KeyError, TypeError): 

126 return None 

127 

128 

129@contextmanager 

130def write_lock( 

131 lancedb_dir: Path | None = None, timeout: float = LOCK_TIMEOUT 

132) -> Generator[None, None, None]: 

133 """Acquire the cross-process file lock then the in-process mutex. 

134 

135 The file lock lives next to the store's data, so cross-process writers 

136 coordinate only when they lock the *same* directory: callers pass their 

137 store's ``lancedb_dir`` (a per-instance ``Lilbee`` uses its own dir). 

138 ``None`` falls back to the global ``cfg.lancedb_dir``. 

139 

140 The two stages share one budget: the time spent waiting on the file lock is 

141 deducted before waiting on the mutex (plus a small ``_MUTEX_MIN_WAIT`` floor), 

142 so a 30s request cannot stall for roughly twice that. 

143 """ 

144 deadline = time.monotonic() + timeout 

145 lock_path = _lock_path(lancedb_dir) 

146 # The first write to a per-instance store can run before its data dir exists; 

147 # the file lock cannot be created in a missing directory. 

148 lock_path.parent.mkdir(parents=True, exist_ok=True) 

149 flock = FileLock(lock_path) 

150 try: 

151 flock.acquire(timeout=timeout) 

152 except FileLockTimeout: 

153 raise LockTimeoutError("Timed out waiting for exclusive file lock") from None 

154 try: 

155 # Floor the mutex budget so a file lock that wins right at the deadline 

156 # still gets a brief blocking attempt instead of a zero-timeout poll that 

157 # spuriously fails when another thread holds the mutex for an instant. 

158 remaining = max(_MUTEX_MIN_WAIT, deadline - time.monotonic()) 

159 acquired = _write_mutex.acquire(timeout=remaining) 

160 if not acquired: 

161 raise LockTimeoutError("Timed out waiting for write lock") 

162 try: 

163 yield 

164 finally: 

165 _write_mutex.release() 

166 finally: 

167 flock.release()