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

1"""Per-process ingest lock registry. 

2 

3A runtime concurrency primitive shared by the HTTP add-files handler, the 

4TUI ingest task, and any other surface that wants to serialize concurrent 

5ingest of the same source file. Lives at the runtime layer so callers in 

6core/server/cli/tui can all use it without dragging in HTTP-layer code. 

7""" 

8 

9from __future__ import annotations 

10 

11import asyncio 

12from pathlib import Path 

13 

14 

15class IngestLockRegistry: 

16 """Per-source ingest locks with a serialized check-and-acquire step. 

17 

18 The registry lock serializes lock creation and the check-and-acquire 

19 so concurrent ``/api/add`` calls cannot TOCTOU between 

20 ``locked()`` and ``acquire()``. One instance is held by ``Services`` 

21 and discarded by ``reset_services()``. 

22 """ 

23 

24 def __init__(self) -> None: 

25 self._locks: dict[str, asyncio.Lock] = {} 

26 self._registry_lock: asyncio.Lock | None = None 

27 

28 def _get_registry_lock(self) -> asyncio.Lock: 

29 if self._registry_lock is None: 

30 self._registry_lock = asyncio.Lock() 

31 return self._registry_lock 

32 

33 def reset(self) -> None: 

34 """Test hook: clear per-source locks and the registry lock.""" 

35 self._locks.clear() 

36 self._registry_lock = None 

37 

38 async def try_acquire(self, name: str) -> asyncio.Lock | None: 

39 """Acquire the lock for ``name`` or return ``None`` if already held.""" 

40 async with self._get_registry_lock(): 

41 lock = self._locks.get(name) 

42 if lock is None: 

43 lock = asyncio.Lock() 

44 self._locks[name] = lock 

45 if lock.locked(): 

46 return None 

47 await lock.acquire() 

48 return lock 

49 

50 @staticmethod 

51 def canonical_source_name(p_str: str) -> str: 

52 """Basename of *p_str*, the label a registered source root keys under. 

53 

54 ``register_sources`` records a server-side path by its basename, so 

55 /api/add locks on that. Uploads keep their relative layout and pass the 

56 relative path instead, which is why the caller picks the key rather than 

57 this class. 

58 """ 

59 return Path(p_str).name 

60 

61 async def acquire(self, names: list[str]) -> tuple[list[tuple[str, asyncio.Lock]], list[str]]: 

62 """Return ``(acquired, busy)`` partitioning of ``names`` by lock state. 

63 

64 Each name is the source identity as it will exist on disk; the caller 

65 derives it, because how a path maps to a stored source differs per 

66 ingest surface. 

67 """ 

68 acquired: list[tuple[str, asyncio.Lock]] = [] 

69 busy: list[str] = [] 

70 seen: set[str] = set() 

71 for name in names: 

72 if name in seen: 

73 continue 

74 seen.add(name) 

75 lock = await self.try_acquire(name) 

76 if lock is None: 

77 busy.append(name) 

78 else: 

79 acquired.append((name, lock)) 

80 return acquired, busy 

81 

82 def release(self, acquired: list[tuple[str, asyncio.Lock]]) -> None: 

83 """Release every lock in ``acquired`` and evict its registry entry. 

84 

85 Runs synchronously (no ``await``), so it is atomic with respect to 

86 ``try_acquire`` on the event loop. ``try_acquire`` only ever acquires a 

87 free lock, so these per-source locks never accrue waiters; dropping the 

88 entry once released keeps the registry from growing one lock per distinct 

89 filename for the whole process lifetime. Safe to call multiple times. 

90 """ 

91 while acquired: 

92 name, lock = acquired.pop() 

93 if lock.locked(): 

94 lock.release() 

95 if self._locks.get(name) is lock: 

96 del self._locks[name]