Coverage for src/evutils/io/_prefetch.py: 96%

1"""Background-thread prefetching for :class:`~evutils.io.EventReader`. 

2 

3The reader's windowing loop runs in a worker thread and pushes finished 

4windows into a small bounded queue, so decoding the next window overlaps the 

5caller's processing of the current one. This works because 

6 

7* the native parsers are called through ctypes, which releases the GIL for 

8 the duration of the C call (measured: a Python thread keeps ~90% of its 

9 solo throughput while a decode runs), and 

10* every yielded window is already an independent copy of the reader's reused 

11 accumulator storage, so the worker can never mutate chunks the consumer 

12 still holds. 

13 

14When it helps: any pipeline where per-chunk processing takes meaningful time 

15-- numpy transforms, writing elsewhere, GPU inference (where the read becomes 

16entirely free). When it does not: processing that already saturates memory 

17bandwidth can get *slower* with prefetching enabled, because decode competes 

18for the same bandwidth. Hence prefetching is opt-in 

19(``EventReader(..., async_read=True)``). 

20""" 

21from __future__ import annotations 

22 

23import queue 

24import threading 

25from typing import Any, Callable, Iterator 

26 

27#: How many finished windows may be buffered ahead of the consumer. Depth 2 is 

28#: enough to decouple producer and consumer; larger values only cost memory 

29#: (depth x window size). 

30DEFAULT_DEPTH = 2 

31 

32 

33class PrefetchIterator: 

34 """Iterate a chunk source through a bounded queue filled by a worker thread. 

35 

36 Semantics match plain iteration: same chunks in the same order, and an 

37 exception raised by the source is re-raised to the consumer at the point 

38 it would have occurred. Use :meth:`close` (or exhaust the iterator, or use 

39 it as a context manager) to release the worker early; an abandoned, 

40 unclosed iterator does not deadlock -- the worker is a daemon thread 

41 parked on a stop-aware put. 

42 

43 Parameters 

44 ---------- 

45 source : Iterator 

46 The synchronous chunk iterator to drain (runs entirely in the worker 

47 thread; it must not be touched by anyone else while this is alive). 

48 depth : int 

49 Maximum number of chunks buffered ahead of the consumer. 

50 on_finish : callable, optional 

51 Called exactly once, from whichever thread finishes the iterator 

52 (exhaustion or :meth:`close`). The EventReader uses it to clear its 

53 active-prefetch guard. 

54 

55 """ 

56 

57 _SENTINEL: Any = object() 

58 

59 def __init__(self, source: Iterator[Any], depth: int = DEFAULT_DEPTH, 

60 on_finish: Callable[[], None] | None = None): 

61 if depth < 1: 

62 raise ValueError("prefetch depth must be >= 1") 

63 self._source = source 

64 self._queue: queue.Queue[Any] = queue.Queue(maxsize=depth) 

65 self._stop = threading.Event() 

66 self._exc: BaseException | None = None 

67 self._finished = False 

68 self._on_finish = on_finish 

69 self._worker = threading.Thread( 

70 target=self._work, name="evutils-prefetch", daemon=True 

71 ) 

72 self._worker.start() 

73 

74 # ------------------------------------------------------------------ # 

75 # Worker side 

76 # ------------------------------------------------------------------ # 

77 def _put(self, item: Any) -> bool: 

78 """Blocking put that stays responsive to :meth:`close`. 

79 

80 Returns False when cancelled. The end-of-stream sentinel must go 

81 through here too: a non-blocking put of the sentinel can be dropped 

82 when the queue is full, leaving the consumer waiting forever. 

83 """ 

84 while not self._stop.is_set(): 

85 try: 

86 self._queue.put(item, timeout=0.1) 

87 return True 

88 except queue.Full: 

89 pass 

90 return False 

91 

92 def _work(self) -> None: 

93 try: 

94 for chunk in self._source: 

95 if not self._put(chunk): 

96 return # cancelled by close() 

97 except BaseException as exc: # propagated to the consumer 

98 self._exc = exc 

99 finally: 

100 close = getattr(self._source, "close", None) 

101 if callable(close): 

102 close() # generators: release reader frame in this thread 

103 self._put(self._SENTINEL) 

104 

105 # ------------------------------------------------------------------ # 

106 # Consumer side 

107 # ------------------------------------------------------------------ # 

108 def __iter__(self) -> "PrefetchIterator": 

109 return self 

110 

111 def __next__(self) -> Any: 

112 if self._finished: 

113 raise StopIteration 

114 item = self._queue.get() 

115 if item is self._SENTINEL: 

116 self._worker.join() 

117 self._finish() 

118 if self._exc is not None: 

119 raise self._exc 

120 raise StopIteration 

121 return item 

122 

123 def close(self) -> None: 

124 """Cancel the worker and drop any buffered chunks. Idempotent.""" 

125 if self._finished: 

126 return 

127 self._stop.set() 

128 # Drain so a worker blocked on put() can observe the stop event. 

129 while True: 

130 try: 

131 self._queue.get_nowait() 

132 except queue.Empty: 

133 break 

134 self._worker.join(timeout=5.0) 

135 self._finish() 

136 

137 def _finish(self) -> None: 

138 if not self._finished: 

139 self._finished = True 

140 if self._on_finish is not None: 

141 self._on_finish() 

142 

143 def __enter__(self) -> "PrefetchIterator": 

144 return self 

145 

146 def __exit__(self, *exc: Any) -> None: 

147 self.close()