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