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 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 

32class PrefetchIterator: 

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

34 

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

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

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

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

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

40 parked on a stop-aware put. 

41 

42 Parameters 

43 ---------- 

44 source : Iterator 

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

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

47 depth : int 

48 Maximum number of chunks buffered ahead of the consumer. 

49 on_finish : callable, optional 

50 Called exactly once, from whichever thread finishes the iterator 

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

52 active-prefetch guard. 

53 

54 """ 

55 

56 _SENTINEL: object = object() 

57 

58 def __init__(self, source: "Iterator[EventArray]", depth: int = DEFAULT_DEPTH, 

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

60 if depth < 1: 

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

62 self._source = source 

63 self._queue: "queue.Queue[EventArray]" = queue.Queue(maxsize=depth) 

64 self._stop = threading.Event() 

65 self._exc: BaseException | None = None 

66 self._finished = False 

67 self._on_finish = on_finish 

68 self._worker = threading.Thread( 

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

70 ) 

71 self._worker.start() 

72 

73 # ------------------------------------------------------------------ # 

74 # Worker side 

75 # ------------------------------------------------------------------ # 

76 def _put(self, item: "EventArray") -> bool: 

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

78 

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

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

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

82 """ 

83 while not self._stop.is_set(): 

84 try: 

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

86 return True 

87 except queue.Full: 

88 pass 

89 return False 

90 

91 def _work(self) -> None: 

92 try: 

93 for chunk in self._source: 

94 if not self._put(chunk): 

95 return # cancelled by close() 

96 except BaseException as exc: # propagated to the consumer 

97 self._exc = exc 

98 finally: 

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

100 if callable(close): 

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

102 self._put(self._SENTINEL) 

103 

104 # ------------------------------------------------------------------ # 

105 # Consumer side 

106 # ------------------------------------------------------------------ # 

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

108 return self 

109 

110 def __next__(self) -> "EventArray": 

111 if self._finished: 

112 raise StopIteration 

113 item = self._queue.get() 

114 if item is self._SENTINEL: 

115 self._worker.join() 

116 self._finish() 

117 if self._exc is not None: 

118 raise self._exc 

119 raise StopIteration 

120 return item 

121 

122 def close(self) -> None: 

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

124 if self._finished: 

125 return 

126 self._stop.set() 

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

128 while True: 

129 try: 

130 self._queue.get_nowait() 

131 except queue.Empty: 

132 break 

133 self._worker.join(timeout=5.0) 

134 self._finish() 

135 

136 def _finish(self) -> None: 

137 if not self._finished: 

138 self._finished = True 

139 if self._on_finish is not None: 

140 self._on_finish() 

141 

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

143 return self 

144 

145 def __exit__(self, *exc: object) -> None: 

146 self.close()