hat.event.backends.lmdb.backend

  1from collections.abc import Collection
  2from pathlib import Path
  3import asyncio
  4import collections
  5import contextlib
  6import logging
  7import typing
  8
  9from hat import aio
 10from hat import json
 11
 12from hat.event.backends.lmdb import common
 13from hat.event.backends.lmdb import environment
 14from hat.event.backends.lmdb import latestdb
 15from hat.event.backends.lmdb import refdb
 16from hat.event.backends.lmdb import systemdb
 17from hat.event.backends.lmdb import timeseriesdb
 18from hat.event.backends.lmdb.conditions import Conditions
 19
 20
 21mlog = logging.getLogger(__name__)
 22
 23cleanup_max_results = 1024
 24
 25flush_queue_size = 4096
 26
 27max_registered_count = 256 * 1024
 28
 29default_timeseries_max_results = 4096
 30
 31default_timeseries_event_type_cache_size = 256 * 1024
 32
 33version = '0.9'
 34
 35
 36class Databases(typing.NamedTuple):
 37    system: systemdb.SystemDb
 38    latest: latestdb.LatestDb
 39    timeseries: timeseriesdb.TimeseriesDb
 40    ref: refdb.RefDb
 41
 42
 43class Changes(typing.NamedTuple):
 44    system: systemdb.Changes
 45    latest: latestdb.Changes
 46    timeseries: timeseriesdb.Changes
 47    ref: refdb.Changes
 48
 49
 50async def create(conf: json.Data,
 51                 registered_events_cb: common.BackendRegisteredEventsCb | None,
 52                 flushed_events_cb: common.BackendFlushedEventsCb | None
 53                 ) -> 'LmdbBackend':
 54    backend = LmdbBackend()
 55    backend._registered_events_cb = registered_events_cb
 56    backend._flushed_events_cb = flushed_events_cb
 57    backend._conditions = Conditions(conf['conditions'])
 58    backend._loop = asyncio.get_running_loop()
 59    backend._flush_queue = aio.Queue(flush_queue_size)
 60    backend._registered_count = 0
 61    backend._registered_queue = collections.deque()
 62    backend._async_group = aio.Group()
 63
 64    backend._env = await environment.create(Path(conf['db_path']))
 65    backend.async_group.spawn(aio.call_on_done, backend._env.wait_closing(),
 66                              backend.close)
 67
 68    try:
 69        latest_subscription = common.create_subscription(
 70            tuple(i) for i in conf['latest']['subscriptions'])
 71
 72        timeseries_partitions = (
 73            timeseriesdb.Partition(
 74                order_by=common.OrderBy[i['order_by']],
 75                subscription=common.create_subscription(
 76                    tuple(event_type) for event_type in i['subscriptions']),
 77                limit=(
 78                    timeseriesdb.Limit(
 79                        min_entries=i['limit'].get('min_entries'),
 80                        max_entries=i['limit'].get('max_entries'),
 81                        duration=i['limit'].get('duration'),
 82                        size=i['limit'].get('size'))
 83                    if 'limit' in i else None))
 84            for i in conf['timeseries'])
 85
 86        timeseries_max_results = conf.get('timeseries_max_results',
 87                                          default_timeseries_max_results)
 88        timeseries_event_type_cache_size = conf.get(
 89            'timeseries_event_type_cache_size',
 90            default_timeseries_event_type_cache_size)
 91
 92        backend._dbs = await backend._env.execute(
 93            _ext_create_dbs, backend._env, conf['identifier'],
 94            backend._conditions, latest_subscription, timeseries_partitions,
 95            timeseries_max_results, timeseries_event_type_cache_size,
 96            backend.flush)
 97
 98        backend.async_group.spawn(backend._flush_loop, conf['flush_period'])
 99        backend.async_group.spawn(backend._cleanup_loop,
100                                  conf['cleanup_period'])
101
102    except BaseException:
103        await aio.uncancellable(backend._env.async_close())
104        raise
105
106    return backend
107
108
109def _ext_create_dbs(env, identifier, conditions, latest_subscription,
110                    timeseries_partitions, timeseries_max_results,
111                    timeseries_event_type_cache_size, flush_cb):
112    with env.ext_begin(write=True) as txn:
113        system_db = systemdb.ext_create(env, txn, version, identifier)
114        latest_db = latestdb.ext_create(env, txn, conditions,
115                                        latest_subscription)
116        timeseries_db = timeseriesdb.ext_create(
117            env, txn, conditions, timeseries_partitions,
118            timeseries_max_results, timeseries_event_type_cache_size,
119            flush_cb)
120        ref_db = refdb.RefDb(env)
121
122    return Databases(system=system_db,
123                     latest=latest_db,
124                     timeseries=timeseries_db,
125                     ref=ref_db)
126
127
128class LmdbBackend(common.Backend):
129
130    @property
131    def async_group(self) -> aio.Group:
132        return self._async_group
133
134    async def get_last_event_id(self,
135                                server_id: int
136                                ) -> common.EventId:
137        if not self.is_open:
138            raise common.BackendClosedError()
139
140        return self._dbs.system.get_last_event_id(server_id)
141
142    async def register(self,
143                       events: Collection[common.Event]
144                       ) -> Collection[common.Event] | None:
145        if not self.is_open:
146            raise common.BackendClosedError()
147
148        for event in events:
149            server_id = event.id.server
150
151            last_event_id = self._dbs.system.get_last_event_id(server_id)
152            last_timestamp = self._dbs.system.get_last_timestamp(server_id)
153
154            if last_event_id >= event.id:
155                mlog.warning("event registration skipped: invalid event id")
156                continue
157
158            if last_timestamp > event.timestamp:
159                mlog.warning("event registration skipped: invalid timestamp")
160                continue
161
162            if not self._conditions.matches(event):
163                mlog.warning(
164                    "event %s registration skipped: invalid conditions",
165                    event.type)
166                continue
167
168            refs = collections.deque()
169
170            latest_result = self._dbs.latest.add(event)
171
172            if latest_result.added_ref:
173                refs.append(latest_result.added_ref)
174
175            if latest_result.removed_ref:
176                self._dbs.ref.remove(*latest_result.removed_ref)
177
178            refs.extend(self._dbs.timeseries.add(event))
179
180            if not refs:
181                continue
182
183            self._dbs.ref.add(event, refs)
184            self._dbs.system.set_last_event_id(event.id)
185            self._dbs.system.set_last_timestamp(server_id, event.timestamp)
186
187        self._registered_queue.append(events)
188        self._registered_count += len(events)
189
190        if self._registered_count > max_registered_count:
191            await self._flush_queue.put(None)
192
193        if self._registered_events_cb:
194            await aio.call(self._registered_events_cb, events)
195
196        return events
197
198    async def query(self,
199                    params: common.QueryParams
200                    ) -> common.QueryResult:
201        if not self.is_open:
202            raise common.BackendClosedError()
203
204        if isinstance(params, common.QueryLatestParams):
205            return self._dbs.latest.query(params)
206
207        if isinstance(params, common.QueryTimeseriesParams):
208            return await self._dbs.timeseries.query(params)
209
210        if isinstance(params, common.QueryServerParams):
211            return await self._dbs.ref.query(params)
212
213        raise ValueError('unsupported params type')
214
215    async def flush(self):
216        try:
217            future = self._loop.create_future()
218            await self._flush_queue.put(future)
219            await future
220
221        except aio.QueueClosedError:
222            raise common.BackendClosedError()
223
224    async def _flush_loop(self, flush_period):
225        futures = collections.deque()
226
227        async def cleanup():
228            with contextlib.suppress(Exception):
229                await self._flush()
230
231            await self._env.async_close()
232
233        try:
234            while True:
235                try:
236                    future = await aio.wait_for(self._flush_queue.get(),
237                                                flush_period)
238                    futures.append(future)
239
240                except asyncio.TimeoutError:
241                    pass
242
243                except aio.CancelledWithResultError as e:
244                    if e.result:
245                        futures.append(e.result)
246
247                    raise
248
249                while not self._flush_queue.empty():
250                    futures.append(self._flush_queue.get_nowait())
251
252                await aio.uncancellable(self._flush())
253
254                while futures:
255                    future = futures.popleft()
256                    if future and not future.done():
257                        future.set_result(None)
258
259        except Exception as e:
260            mlog.error('backend flush error: %s', e, exc_info=e)
261
262        finally:
263            self.close()
264            self._flush_queue.close()
265
266            while not self._flush_queue.empty():
267                futures.append(self._flush_queue.get_nowait())
268
269            for future in futures:
270                if future and not future.done():
271                    future.set_exception(common.BackendClosedError())
272
273            await aio.uncancellable(cleanup())
274
275    async def _cleanup_loop(self, cleanup_period):
276        try:
277            while True:
278                await asyncio.sleep(0)
279
280                repeat = await self._env.execute(_ext_cleanup, self._env,
281                                                 self._dbs, common.now())
282                if repeat:
283                    continue
284
285                await asyncio.sleep(cleanup_period)
286
287        except Exception as e:
288            mlog.error('backend cleanup error: %s', e, exc_info=e)
289
290        finally:
291            self.close()
292
293    async def _flush(self):
294        if not self._env.is_open:
295            return
296
297        self._registered_count = 0
298        registered_queue, self._registered_queue = (self._registered_queue,
299                                                    collections.deque())
300
301        changes = Changes(system=self._dbs.system.create_changes(),
302                          latest=self._dbs.latest.create_changes(),
303                          timeseries=self._dbs.timeseries.create_changes(),
304                          ref=self._dbs.ref.create_changes())
305
306        # TODO lock period between create_changes and locking executor
307        #      (timeseries and ref must write changes before new queries are
308        #      allowed)
309
310        await self._env.execute(_ext_flush, self._env, self._dbs, changes)
311
312        if not self._flushed_events_cb:
313            return
314
315        while registered_queue:
316            events = registered_queue.popleft()
317            await aio.call(self._flushed_events_cb, events)
318
319
320def _ext_flush(env, dbs, changes):
321    with env.ext_begin(write=True) as txn:
322        dbs.system.ext_write(txn, changes.system)
323        dbs.latest.ext_write(txn, changes.latest)
324        dbs.timeseries.ext_write(txn, changes.timeseries)
325        dbs.ref.ext_write(txn, changes.ref)
326
327
328def _ext_cleanup(env, dbs, now):
329    with env.ext_begin(write=True) as txn:
330        result = dbs.timeseries.ext_cleanup(txn, now, cleanup_max_results)
331        if not result:
332            return False
333
334        dbs.ref.ext_cleanup(txn, result)
335
336    return len(result) >= cleanup_max_results