Coverage for arrakis/mux.py: 73.1%

1# Copyright (c) 2022, California Institute of Technology and contributors

2#

3# You should have received a copy of the licensing terms for this

4# software included in the file "LICENSE" located in the top-level

5# directory of this package. If you did not, you can view a copy at

6# https://git.ligo.org/ngdd/arrakis-server/-/raw/main/LICENSE

7

8 import heapq

9 import warnings

10 from collections import defaultdict

11 from collections.abc import Iterable, Iterator, Mapping

12 from dataclasses import dataclass

13 from datetime import timedelta

14 from enum import Enum, auto

15 from typing import Generic, TypeVar

16

17 import gpstime

18

19 from .block import Time

20

21 T = TypeVar("T")

22

23 DEFAULT_TIMEOUT = timedelta(seconds=1)

24

25

26 class OnDrop(Enum):

27 IGNORE = auto()

28 RAISE = auto()

29 WARN = auto()

30

31

32@dataclass

33 class MuxedData(Mapping, Generic[T]):

34 """Container that holds timestamped data.

35

36 Parameters

37 ----------

38 time : int

39 The timestamp associated with this data, in nanoseconds.

40 data : dict[str, T]

41 The keyed data.

42

43 """

44

45 time: int

46 data: dict[str, T]

47

48 def __getitem__(self, index: str) -> T:

49 return self.data[index]

50

51 def __iter__(self) -> Iterator[str]:

52 return iter(self.data)

53

54 def __len__(self) -> int:

55 return len(self.data)

56

57

58 class Muxer(Generic[T]):

59 """A data structure that multiplexes items from multiple named streams.

60

61 Given items from multiple named streams with monotonically increasing

62 integer timestamps, this data structure can be used to pull out sets of

63 synchronized items (items all with the same timestamp).

64

65 The oldest items will be held until either all named streams are

66 available or until the timeout has been reached. If a start time has been

67 set, any items with an older timestamp will be rejected.

68

69 Parameters

70 ----------

71 keys : Iterable[str]

72 Identifiers for the named streams to expect when adding items.

73 start : int, optional

74 The GPS time to start muxing items for.

75 If not set, accept items from any time.

76 timeout : timedelta or None, optional

77 The maximum time to wait for messages from named streams, in seconds,

78 before multiplexing. If None is specified, wait indefinitely. Default

79 is 1 second.

80

81 """

82

83 def __init__(

84 self,

85 keys: Iterable[str],

86 start: int | None = None,

87 timeout: timedelta | None = DEFAULT_TIMEOUT,

88 ) -> None:

89 self._keys = set(keys)

90 self._items: dict[int, dict[str, T]] = defaultdict(lambda: defaultdict())

91 self._times: list[int] = []

92 self._last_time = None

93 self._start = start

94 self._last_time = start

95 self._timeout = timeout

96

97 # track when processing started to handle lookback properly

98 self._processing_start_time = int(gpstime.gpsnow() * Time.SECONDS)

99

100 def push(self, time: int, key: str, item: T, on_drop: str = "warn") -> None:

101 """Push an item into the muxer.

102

103 Parameters

104 ----------

105 time : int

106 The timestamp associated with this item.

107 key : str

108 The key stream associated with this item. Must match a key provided

109 at initialization.

110 item : T

111 The item to add.

112 on_drop : str, optional

113 Specifies behavior when the item would be dropped from the muxer,

114 in the case that it was not provided to the muxer before the

115 specified timeout. Options are 'ignore', 'raise', or 'warn'.

116 Default is 'warn'.

117

118 """

119 if key not in self._keys:

120 msg = f"{key} doesn't match keys provided at initialization"

121 raise KeyError(msg)

122

123 if not self._last_time:

124 self._last_time = time

125

126 # skip over items that have already been pulled

127 if time < self._last_time:

128 if self._start is not None and time < self._start:

129 return

130 msg = f"item's timestamp is too old: ({time} < {self._last_time})"

131 match OnDrop[on_drop.lower()]:

132 case OnDrop.IGNORE:

133 return

134 case OnDrop.RAISE:

135 raise ValueError(msg)

136 case OnDrop.WARN:

137 warnings.warn(msg, stacklevel=2)

138

139 # add item

140 if time in self._items:

141 if key not in self._items[time]:

142 self._items[time][key] = item

143 else:

144 heapq.heappush(self._times, time)

145 self._items[time][key] = item

146

147 def pull(self) -> Iterator[MuxedData[T]]:

148 """Pull monotonically increasing synchronized items from the muxer.

149

150 Yields

151 ------

152 MuxedData[T]

153 Synchronized items with a common timestamp, keyed by stream keys.

154

155 """

156 if not self._times:

157 return

158

159 # yield items in monotonically increasing order as long

160 # as conditions are met

161 time = self._times[0]

162 while self._has_all_items(time) or self._are_items_stale(time):

163 yield MuxedData(time, self._items.pop(time))

164 self._last_time = heapq.heappop(self._times)

165 if not self._times:

166 break

167 time = self._times[0]

168

169 def _has_all_items(self, time: int):

170 """Check if a timestamp has all items requested."""

171 return len(self._items[time]) == len(self._keys)

172

173 def _are_items_stale(self, time):

174 """Check if a timestamp is older than the latency cutoff."""

175 if self._timeout is None:

176 return False

177

178 time_now = gpstime.gpsnow()

179 dt_lookback = max(self._processing_start_time - time, 0) / float(Time.SECONDS)

180 dt_timeout = self._timeout.total_seconds()

181 oldest_time_allowed = time_now - dt_lookback - dt_timeout

182 return time <= int(oldest_time_allowed * Time.SECONDS)