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
 97        backend.async_group.spawn(backend._flush_loop, conf['flush_period'])
 98        backend.async_group.spawn(backend._cleanup_loop,
 99                                  conf['cleanup_period'])
100
101    except BaseException:
102        await aio.uncancellable(backend._env.async_close())
103        raise
104
105    return backend
106
107
108def _ext_create_dbs(env, identifier, conditions, latest_subscription,
109                    timeseries_partitions, timeseries_max_results,
110                    timeseries_event_type_cache_size):
111    with env.ext_begin(write=True) as txn:
112        system_db = systemdb.ext_create(env, txn, version, identifier)
113        latest_db = latestdb.ext_create(env, txn, conditions,
114                                        latest_subscription)
115        timeseries_db = timeseriesdb.ext_create(
116            env, txn, conditions, timeseries_partitions,
117            timeseries_max_results, timeseries_event_type_cache_size)
118        ref_db = refdb.RefDb(env)
119
120    return Databases(system=system_db,
121                     latest=latest_db,
122                     timeseries=timeseries_db,
123                     ref=ref_db)
124
125
126class LmdbBackend(common.Backend):
127
128    @property
129    def async_group(self) -> aio.Group:
130        return self._async_group
131
132    async def get_last_event_id(self,
133                                server_id: int
134                                ) -> common.EventId:
135        if not self.is_open:
136            raise common.BackendClosedError()
137
138        return self._dbs.system.get_last_event_id(server_id)
139
140    async def register(self,
141                       events: Collection[common.Event]
142                       ) -> Collection[common.Event] | None:
143        if not self.is_open:
144            raise common.BackendClosedError()
145
146        for event in events:
147            server_id = event.id.server
148
149            last_event_id = self._dbs.system.get_last_event_id(server_id)
150            last_timestamp = self._dbs.system.get_last_timestamp(server_id)
151
152            if last_event_id >= event.id:
153                mlog.warning("event registration skipped: invalid event id")
154                continue
155
156            if last_timestamp > event.timestamp:
157                mlog.warning("event registration skipped: invalid timestamp")
158                continue
159
160            if not self._conditions.matches(event):
161                mlog.warning("event registration skipped: invalid conditions")
162                continue
163
164            refs = collections.deque()
165
166            latest_result = self._dbs.latest.add(event)
167
168            if latest_result.added_ref:
169                refs.append(latest_result.added_ref)
170
171            if latest_result.removed_ref:
172                self._dbs.ref.remove(*latest_result.removed_ref)
173
174            refs.extend(self._dbs.timeseries.add(event))
175
176            if not refs:
177                continue
178
179            self._dbs.ref.add(event, refs)
180            self._dbs.system.set_last_event_id(event.id)
181            self._dbs.system.set_last_timestamp(server_id, event.timestamp)
182
183        self._registered_queue.append(events)
184        self._registered_count += len(events)
185
186        if self._registered_count > max_registered_count:
187            await self._flush_queue.put(None)
188
189        if self._registered_events_cb:
190            await aio.call(self._registered_events_cb, events)
191
192        return events
193
194    async def query(self,
195                    params: common.QueryParams
196                    ) -> common.QueryResult:
197        if not self.is_open:
198            raise common.BackendClosedError()
199
200        if isinstance(params, common.QueryLatestParams):
201            return self._dbs.latest.query(params)
202
203        if isinstance(params, common.QueryTimeseriesParams):
204            return await self._dbs.timeseries.query(params)
205
206        if isinstance(params, common.QueryServerParams):
207            return await self._dbs.ref.query(params)
208
209        raise ValueError('unsupported params type')
210
211    async def flush(self):
212        try:
213            future = self._loop.create_future()
214            await self._flush_queue.put(future)
215            await future
216
217        except aio.QueueClosedError:
218            raise common.BackendClosedError()
219
220    async def _flush_loop(self, flush_period):
221        futures = collections.deque()
222
223        async def cleanup():
224            with contextlib.suppress(Exception):
225                await self._flush()
226
227            await self._env.async_close()
228
229        try:
230            while True:
231                try:
232                    future = await aio.wait_for(self._flush_queue.get(),
233                                                flush_period)
234                    futures.append(future)
235
236                except asyncio.TimeoutError:
237                    pass
238
239                except aio.CancelledWithResultError as e:
240                    if e.result:
241                        futures.append(e.result)
242
243                    raise
244
245                while not self._flush_queue.empty():
246                    futures.append(self._flush_queue.get_nowait())
247
248                await aio.uncancellable(self._flush())
249
250                while futures:
251                    future = futures.popleft()
252                    if future and not future.done():
253                        future.set_result(None)
254
255        except Exception as e:
256            mlog.error('backend flush error: %s', e, exc_info=e)
257
258        finally:
259            self.close()
260            self._flush_queue.close()
261
262            while not self._flush_queue.empty():
263                futures.append(self._flush_queue.get_nowait())
264
265            for future in futures:
266                if future and not future.done():
267                    future.set_exception(common.BackendClosedError())
268
269            await aio.uncancellable(cleanup())
270
271    async def _cleanup_loop(self, cleanup_period):
272        try:
273            while True:
274                await asyncio.sleep(0)
275
276                repeat = await self._env.execute(_ext_cleanup, self._env,
277                                                 self._dbs, common.now())
278                if repeat:
279                    continue
280
281                await asyncio.sleep(cleanup_period)
282
283        except Exception as e:
284            mlog.error('backend cleanup error: %s', e, exc_info=e)
285
286        finally:
287            self.close()
288
289    async def _flush(self):
290        if not self._env.is_open:
291            return
292
293        self._registered_count = 0
294        registered_queue, self._registered_queue = (self._registered_queue,
295                                                    collections.deque())
296
297        changes = Changes(system=self._dbs.system.create_changes(),
298                          latest=self._dbs.latest.create_changes(),
299                          timeseries=self._dbs.timeseries.create_changes(),
300                          ref=self._dbs.ref.create_changes())
301
302        # TODO lock period between create_changes and locking executor
303        #      (timeseries and ref must write changes before new queries are
304        #      allowed)
305
306        await self._env.execute(_ext_flush, self._env, self._dbs, changes)
307
308        if not self._flushed_events_cb:
309            return
310
311        while registered_queue:
312            events = registered_queue.popleft()
313            await aio.call(self._flushed_events_cb, events)
314
315
316def _ext_flush(env, dbs, changes):
317    with env.ext_begin(write=True) as txn:
318        dbs.system.ext_write(txn, changes.system)
319        dbs.latest.ext_write(txn, changes.latest)
320        dbs.timeseries.ext_write(txn, changes.timeseries)
321        dbs.ref.ext_write(txn, changes.ref)
322
323
324def _ext_cleanup(env, dbs, now):
325    with env.ext_begin(write=True) as txn:
326        result = dbs.timeseries.ext_cleanup(txn, now, cleanup_max_results)
327        if not result:
328            return False
329
330        dbs.ref.ext_cleanup(txn, result)
331
332    return len(result) >= cleanup_max_results