Coverage for src / ocarina / infra / drivers_pool.py: 17.28%

1"""Generic driver pool with concurrency control via semaphore. 

2 

3Drivers are never reused — each acquire() disposes after use (clean state). 

4warmup() pre-fills the pool; acquire() creates on demand if pool is empty. 

5 

6Safety guarantees: 

7- No semaphore leaks 

8- Warmup cannot stall silently 

9- Stuck warmup is detected via progress watchdog 

10""" 

11 

12import threading 

13import time 

14from contextlib import contextmanager, suppress 

15from queue import Empty, Queue 

16from threading import Semaphore 

17from typing import TYPE_CHECKING, final 

18 

19if TYPE_CHECKING: 

20 from collections.abc import Iterator 

21 

22 from ocarina.custom_types.built_web_driver import BuiltWebDriver 

23 from ocarina.custom_types.thunk import Thunk 

24 

25 

26class WarmupTimeoutError(Exception): 

27 """Raised when warmup is blocked (mostly some macOS edge cases).""" 

28 

29 

30@final 

31class WebDriversPool[Driver]: 

32 """Thread-safe driver pool. Max concurrency = max_size (semaphore-controlled).""" 

33 

34 def __init__( 

35 self, 

36 create_driver: Thunk[BuiltWebDriver[Driver]], 

37 max_size: int, 

38 warmup_timeout: float | None = None, 

39 ) -> None: 

40 """Initialize the pool. 

41 

42 Args: 

43 create_driver: Factory returning (driver, dispose) tuple. 

44 max_size: Maximum concurrent drivers allowed. 

45 warmup_timeout: Seconds without progress before warmup is aborted. 

46 

47 """ 

48 self._create_driver = create_driver 

49 self._pool: Queue[BuiltWebDriver[Driver]] = Queue(max_size) 

50 self._semaphore = Semaphore(max_size) 

51 self._warmup_timeout = ( 

52 warmup_timeout 

53 if warmup_timeout is not None and warmup_timeout > 0.1 # noqa: PLR2004 

54 else 60.0 * 5 

55 ) 

56 

57 @contextmanager 

58 def acquire(self) -> Iterator[Driver]: 

59 """Yield a driver. Blocks if max_size drivers are in use.""" 

60 try: 

61 driver, dispose = self._pool.get_nowait() 

62 except Empty: 

63 self._semaphore.acquire() 

64 try: 

65 driver, dispose = self._create_driver() 

66 except Exception: 

67 self._semaphore.release() 

68 raise 

69 

70 try: 

71 yield driver 

72 finally: 

73 with suppress(Exception): 

74 dispose() 

75 self._semaphore.release() 

76 

77 def warmup(self) -> None: 

78 """Fully initialize the driver pool, avoiding cold-start latency. 

79 

80 Raises: 

81 WarmupTimeoutError: If no progress is detected. 

82 

83 """ 

84 progress = {"count": 0} 

85 lock = threading.Lock() 

86 stop_event = threading.Event() 

87 

88 def worker() -> None: 

89 while not self._pool.full() and not stop_event.is_set(): 

90 if not self._semaphore.acquire(blocking=False): 

91 break 

92 

93 try: 

94 driver, dispose = self._create_driver() 

95 self._pool.put((driver, dispose)) 

96 

97 with lock: 

98 progress["count"] += 1 

99 

100 except Exception: # noqa: BLE001 

101 self._semaphore.release() 

102 break 

103 

104 t = threading.Thread(target=worker, daemon=True) 

105 t.start() 

106 

107 start_global = time.monotonic() 

108 last_progress = 0 

109 

110 while t.is_alive(): 

111 time.sleep(0.5) 

112 

113 with lock: 

114 current = progress["count"] 

115 

116 if current != last_progress: 

117 last_progress = current 

118 start_global = time.monotonic() 

119 

120 if time.monotonic() - start_global > self._warmup_timeout: 

121 stop_event.set() 

122 msg = ( 

123 "Warmup stalled (no progress detected)." 

124 " " 

125 "Some browser processes may still be running." 

126 " " 

127 "Please check your system (Activity Monitor / Task Manager / Dock)" 

128 " " 

129 "and close any remaining browser instances." 

130 ) 

131 self.shutdown() 

132 raise WarmupTimeoutError(msg) 

133 

134 t.join() 

135 

136 def shutdown(self) -> None: 

137 """Dispose pre-created drivers in the pool. Acquired drivers unaffected.""" 

138 while not self._pool.empty(): 

139 _, dispose = self._pool.get_nowait() 

140 with suppress(Exception): 

141 dispose() 

142 self._semaphore.release()