hat.event.backends.lmdb.timeseriesdb

  1from collections.abc import Awaitable, Callable, Iterable
  2import collections
  3import functools
  4import itertools
  5import typing
  6
  7import lmdb
  8
  9from hat.event.backends.lmdb import common
 10from hat.event.backends.lmdb import environment
 11from hat.event.backends.lmdb.conditions import Conditions
 12
 13
 14Changes: typing.TypeAlias = dict[common.PartitionId,
 15                                 collections.deque[tuple[common.Timestamp,
 16                                                         common.Event]]]
 17
 18
 19class Limit(typing.NamedTuple):
 20    min_entries: int | None = None
 21    max_entries: int | None = None
 22    duration: float | None = None
 23    size: int | None = None
 24
 25
 26class Partition(typing.NamedTuple):
 27    order_by: common.OrderBy
 28    subscription: common.Subscription
 29    limit: Limit | None
 30
 31
 32def ext_create(env: environment.Environment,
 33               txn: lmdb.Transaction,
 34               conditions: Conditions,
 35               partitions: Iterable[Partition],
 36               max_results: int,
 37               event_type_cache_size: int,
 38               flush_cb: Callable[[], Awaitable[None]]
 39               ) -> 'TimeseriesDb':
 40    db = TimeseriesDb()
 41    db._env = env
 42    db._conditions = conditions
 43    db._max_results = max_results
 44    db._flush_cb = flush_cb
 45    db._changes = collections.defaultdict(collections.deque)
 46    db._event_type_partitions = {}  # type: dict[common.EventType, Collection[tuple[common.PartitionId, Partition]]]  # NOQA
 47
 48    # depending on dict order
 49    db._partitions = dict(_ext_init_partitions(env, txn, partitions))
 50
 51    partitions_collection = common.create_event_type_collection(
 52        (partition.subscription, (partition_id, partition))
 53        for partition_id, partition in db._partitions.items())
 54
 55    db._get_event_type_partitions = functools.lru_cache(
 56        maxsize=event_type_cache_size)(
 57        lambda event_type: list(partitions_collection.get(event_type)))
 58
 59    return db
 60
 61
 62class TimeseriesDb:
 63
 64    def add(self,
 65            event: common.Event
 66            ) -> Iterable[common.EventRef]:
 67        partitions = self._get_event_type_partitions(event.type)
 68
 69        for partition_id, partition in partitions:
 70            if partition.order_by == common.OrderBy.TIMESTAMP:
 71                timestamp = event.timestamp
 72
 73            elif partition.order_by == common.OrderBy.SOURCE_TIMESTAMP:
 74                if event.source_timestamp is None:
 75                    continue
 76
 77                timestamp = event.source_timestamp
 78
 79            else:
 80                raise ValueError('unsupported order by')
 81
 82            self._changes[partition_id].append((timestamp, event))
 83
 84            yield common.TimeseriesEventRef(
 85                (partition_id, timestamp, event.id))
 86
 87    async def query(self,
 88                    params: common.QueryTimeseriesParams
 89                    ) -> common.QueryResult:
 90        subscription = (common.create_subscription(params.event_types)
 91                        if params.event_types is not None else None)
 92
 93        max_results = (params.max_results
 94                       if params.max_results is not None and
 95                       params.max_results < self._max_results
 96                       else self._max_results)
 97
 98        for partition_id, partition in self._partitions.items():
 99            if partition.order_by != params.order_by:
100                continue
101
102            if (subscription and
103                    subscription.isdisjoint(partition.subscription)):
104                continue
105
106            return await self._query_partition(partition_id, params,
107                                               subscription, max_results)
108
109        return common.QueryResult(events=[],
110                                  more_follows=False)
111
112    def create_changes(self) -> Changes:
113        changes, self._changes = (self._changes,
114                                  collections.defaultdict(collections.deque))
115        return changes
116
117    def ext_write(self,
118                  txn: lmdb.Transaction,
119                  changes: Changes):
120        data = (((partition_id, timestamp, event.id), event)
121                for partition_id, partition_changes in changes.items()
122                for timestamp, event in partition_changes)
123        self._env.ext_write(txn, common.DbType.TIMESERIES_DATA, data)
124
125        counts = ((partition_id, len(partition_changes))
126                  for partition_id, partition_changes in changes.items())
127        _ext_inc_partition_count(self._env, txn, counts)
128
129    def ext_cleanup(self,
130                    txn: lmdb.Transaction,
131                    now: common.Timestamp,
132                    max_results: int | None = None,
133                    ) -> collections.deque[tuple[common.EventId,
134                                                 common.EventRef]]:
135        result = collections.deque()
136
137        for partition_id, partition in self._partitions.items():
138            if not partition.limit:
139                continue
140
141            partition_max_results = (max_results - len(result)
142                                     if max_results is not None else None)
143            if partition_max_results is not None and partition_max_results < 1:
144                break
145
146            result.extend(_ext_cleanup_partition(self._env, txn, now,
147                                                 partition_id, partition.limit,
148                                                 partition_max_results))
149
150        return result
151
152    async def _query_partition(self, partition_id, params, subscription,
153                               max_results):
154        events = collections.deque()
155        changes = self._changes
156
157        filter = _Filter(subscription=subscription,
158                         t_from=params.t_from,
159                         t_to=params.t_to,
160                         source_t_from=params.source_t_from,
161                         source_t_to=params.source_t_to,
162                         max_results=max_results + 1,
163                         last_event_id=params.last_event_id)
164
165        if params.order_by == common.OrderBy.TIMESTAMP:
166            if params.order == common.Order.DESCENDING:
167                events.extend(_query_partition_changes(
168                    changes[partition_id], params, filter))
169
170                if not filter.done:
171                    events.extend(await self._env.execute(
172                        _ext_query_partition_events, self._env,
173                        self._conditions, partition_id, params, filter))
174
175            elif params.order == common.Order.ASCENDING:
176                events.extend(await self._env.execute(
177                    _ext_query_partition_events, self._env, self._conditions,
178                    partition_id, params, filter))
179
180                if not filter.done:
181                    events.extend(_query_partition_changes(
182                        changes[partition_id], params, filter))
183
184            else:
185                raise ValueError('unsupported order')
186
187        elif params.order_by == common.OrderBy.SOURCE_TIMESTAMP:
188            await self._flush_cb()
189
190            events.extend(await self._env.execute(
191                _ext_query_partition_events, self._env, self._conditions,
192                partition_id, params, filter))
193
194        else:
195            raise ValueError('unsupported order by')
196
197        more_follows = len(events) > max_results
198        while len(events) > max_results:
199            events.pop()
200
201        return common.QueryResult(events=events,
202                                  more_follows=more_follows)
203
204
205def _ext_init_partitions(env, txn, partitions):
206
207    def data_to_cache_key(data):
208        return (data['order'],
209                tuple(tuple(event_type)
210                      for event_type in data['subscriptions']))
211
212    max_partition_id = 0
213    cache = {}
214
215    for partition_id, partition_data in env.ext_read(
216            txn, common.DbType.TIMESERIES_PARTITION):
217        if partition_id > max_partition_id:
218            max_partition_id = partition_id
219
220        cache_key = data_to_cache_key(partition_data)
221        cache[cache_key] = partition_id
222
223    next_partition_ids = itertools.count(max_partition_id + 1)
224
225    for partition in partitions:
226        event_types = sorted(partition.subscription.get_query_types())
227        partition_data = {'order': partition.order_by.value,
228                          'subscriptions': [list(i) for i in event_types]}
229
230        cache_key = data_to_cache_key(partition_data)
231        partition_id = cache.get(cache_key)
232
233        if partition_id is None:
234            partition_id = next(next_partition_ids)
235            cache[cache_key] = partition_id
236            env.ext_write(txn, common.DbType.TIMESERIES_PARTITION,
237                          [(partition_id, partition_data)])
238
239        yield partition_id, partition
240
241
242def _ext_query_partition_events(env, conditions, partition_id, params,
243                                filter):
244    if params.order_by == common.OrderBy.TIMESTAMP:
245        events = _ext_query_partition_events_range(
246            env, partition_id, params.t_from, params.t_to, params.order)
247
248    elif params.order_by == common.OrderBy.SOURCE_TIMESTAMP:
249        events = _ext_query_partition_events_range(
250            env, partition_id, params.source_t_from, params.source_t_to,
251            params.order)
252
253    else:
254        raise ValueError('unsupported order by')
255
256    events = (event for event in events if conditions.matches(event))
257    events = filter.process(events)
258    return collections.deque(events)
259
260
261def _ext_query_partition_events_range(env, partition_id, t_from, t_to, order):
262    db_def = common.db_defs[common.DbType.TIMESERIES_DATA]
263
264    if not t_from:
265        t_from = common.min_timestamp
266
267    from_key = partition_id, t_from, common.EventId(0, 0, 0)
268    encoded_from_key = db_def.encode_key(from_key)
269
270    if not t_to:
271        t_to = common.max_timestamp
272
273    to_key = (partition_id,
274              t_to,
275              common.EventId((1 << 64) - 1, (1 << 64) - 1, (1 << 64) - 1))
276    encoded_to_key = db_def.encode_key(to_key)
277
278    with env.ext_begin() as txn:
279        with env.ext_cursor(txn, common.DbType.TIMESERIES_DATA) as cursor:
280            if order == common.Order.DESCENDING:
281                encoded_start_key, encoded_stop_key = (encoded_to_key,
282                                                       encoded_from_key)
283
284                if cursor.set_range(encoded_start_key):
285                    more = cursor.prev()
286                else:
287                    more = cursor.last()
288
289                while more and encoded_stop_key <= bytes(cursor.key()):
290                    yield db_def.decode_value(cursor.value())
291                    more = cursor.prev()
292
293            elif order == common.Order.ASCENDING:
294                encoded_start_key, encoded_stop_key = (encoded_from_key,
295                                                       encoded_to_key)
296
297                more = cursor.set_range(encoded_start_key)
298
299                while more and bytes(cursor.key()) <= encoded_stop_key:
300                    yield db_def.decode_value(cursor.value())
301                    more = cursor.next()
302
303            else:
304                raise ValueError('unsupported order')
305
306
307def _ext_get_partition_count(env, txn, partition_id):
308    db_def = common.db_defs[common.DbType.TIMESERIES_COUNT]
309
310    with env.ext_cursor(txn, common.DbType.TIMESERIES_COUNT) as cursor:
311        encoded_key = db_def.encode_key(partition_id)
312        encoded_value = cursor.get(encoded_key)
313
314        return db_def.decode_value(encoded_value) if encoded_value else 0
315
316
317def _ext_set_partition_count(env, txn, partition_id, count):
318    env.ext_write(txn, common.DbType.TIMESERIES_COUNT, [(partition_id, count)])
319
320
321def _ext_inc_partition_count(env, txn, partition_counts):
322    db_def = common.db_defs[common.DbType.TIMESERIES_COUNT]
323
324    with env.ext_cursor(txn, common.DbType.TIMESERIES_COUNT) as cursor:
325        for partition_id, count in partition_counts:
326            encoded_key = db_def.encode_key(partition_id)
327            encoded_value = cursor.get(encoded_key)
328
329            value = db_def.decode_value(encoded_value) if encoded_value else 0
330            inc_value = value + count
331
332            encoded_value = db_def.encode_value(inc_value)
333            cursor.put(encoded_key, encoded_value)
334
335
336def _ext_cleanup_partition(env, txn, now, partition_id, limit, max_results):
337    db_def = common.db_defs[common.DbType.TIMESERIES_DATA]
338
339    timestamp = common.min_timestamp
340    start_key = (partition_id,
341                 timestamp,
342                 common.EventId(0, 0, 0))
343    stop_key = ((partition_id + 1),
344                timestamp,
345                common.EventId(0, 0, 0))
346
347    encoded_start_key = db_def.encode_key(start_key)
348    encoded_stop_key = db_def.encode_key(stop_key)
349
350    min_entries = limit.min_entries or 0
351    max_entries = None
352    encoded_duration_key = None
353
354    if limit.size is not None:
355        stat = env.ext_stat(txn, common.DbType.TIMESERIES_DATA)
356
357        if stat['entries']:
358            total_size = stat['psize'] * (stat['branch_pages'] +
359                                          stat['leaf_pages'] +
360                                          stat['overflow_pages'])
361            entry_size = total_size / stat['entries']
362            max_entries = int(limit.size / entry_size)
363
364    if limit.max_entries is not None:
365        max_entries = (limit.max_entries if max_entries is None
366                       else min(max_entries, limit.max_entries))
367
368    if limit.duration is not None:
369        duration_key = (partition_id,
370                        now.add(-limit.duration),
371                        common.EventId(0, 0, 0))
372        encoded_duration_key = db_def.encode_key(duration_key)
373
374    result_count = 0
375    entries_count = _ext_get_partition_count(env, txn, partition_id)
376
377    with env.ext_cursor(txn, common.DbType.TIMESERIES_DATA) as cursor:
378        more = cursor.set_range(encoded_start_key)
379        while more:
380            if max_results is not None and result_count >= max_results:
381                break
382
383            if entries_count - result_count <= min_entries:
384                break
385
386            encoded_key = bytes(cursor.key())
387            if encoded_key >= encoded_stop_key:
388                break
389
390            if ((max_entries is None or
391                 entries_count - result_count <= max_entries) and
392                (encoded_duration_key is None or
393                 encoded_key >= encoded_duration_key)):
394                break
395
396            key = db_def.decode_key(encoded_key)
397            event_id = key[2]
398
399            more = cursor.delete()
400            result_count += 1
401
402            yield event_id, common.TimeseriesEventRef(key)
403
404    if result_count > 0:
405        _ext_set_partition_count(env, txn, partition_id,
406                                 entries_count - result_count)
407
408
409def _query_partition_changes(changes, params, filter):
410    assert params.order_by == common.OrderBy.TIMESTAMP
411
412    if params.order == common.Order.DESCENDING:
413        events = (event for _, event in reversed(changes))
414
415        if params.t_to is not None:
416            events = itertools.dropwhile(
417                lambda i: params.t_to < i.timestamp,
418                events)
419
420        if params.t_from is not None:
421            events = itertools.takewhile(
422                lambda i: params.t_from <= i.timestamp,
423                events)
424
425    elif params.order == common.Order.ASCENDING:
426        events = (event for _, event in changes)
427
428        if params.t_from is not None:
429            events = itertools.dropwhile(
430                lambda i: i.timestamp < params.t_from,
431                events)
432
433        if params.t_to is not None:
434            events = itertools.takewhile(
435                lambda i: i.timestamp <= params.t_to,
436                events)
437
438    else:
439        raise ValueError('unsupported order')
440
441    return filter.process(events)
442
443
444class _Filter:
445
446    def __init__(self,
447                 subscription: common.Subscription,
448                 t_from: common.Timestamp | None,
449                 t_to: common.Timestamp | None,
450                 source_t_from: common.Timestamp | None,
451                 source_t_to: common.Timestamp | None,
452                 max_results: int,
453                 last_event_id: common.EventId | None):
454        self._subscription = subscription
455        self._t_from = t_from
456        self._t_to = t_to
457        self._source_t_from = source_t_from
458        self._source_t_to = source_t_to
459        self._max_results = max_results
460        self._last_event_id = last_event_id
461
462    @property
463    def done(self):
464        return self._max_results < 1
465
466    def process(self, events: Iterable[common.Event]):
467        for event in events:
468            if self._max_results < 1:
469                return
470
471            if self._last_event_id:
472                if event.id == self._last_event_id:
473                    self._last_event_id = None
474
475                continue
476
477            if self._t_from is not None and event.timestamp < self._t_from:
478                continue
479
480            if self._t_to is not None and self._t_to < event.timestamp:
481                continue
482
483            if self._source_t_from is not None and (
484                    event.source_timestamp is None or
485                    event.source_timestamp < self._source_t_from):
486                continue
487
488            if self._source_t_to is not None and (
489                    event.source_timestamp is None or
490                    self._source_t_to < event.source_timestamp):
491                continue
492
493            if (self._subscription and
494                    not self._subscription.matches(event.type)):
495                continue
496
497            self._max_results -= 1
498            yield event