lbry-sdk/lbry/wallet/server/db/db.py

import struct
import rocksdb
from typing import Optional
from lbry.wallet.server.db import DB_PREFIXES
from lbry.wallet.server.db.revertable import RevertableOpStack, RevertablePut, RevertableDelete


class RocksDBStore:
    def __init__(self, path: str, cache_mb: int, max_open_files: int, secondary_path: str = ''):
        # Use snappy compression (the default)
        self.path = path
        self._max_open_files = max_open_files
        self.db = rocksdb.DB(path, self.get_options(), secondary_name=secondary_path)
        # self.multi_get = self.db.multi_get

    def get_options(self):
        return rocksdb.Options(
            create_if_missing=True, use_fsync=True, target_file_size_base=33554432,
            max_open_files=self._max_open_files
        )

    def get(self, key: bytes, fill_cache: bool = True) -> Optional[bytes]:
        return self.db.get(key, fill_cache=fill_cache)

    def iterator(self, reverse=False, start=None, stop=None, include_start=True, include_stop=False, prefix=None,
                 include_key=True, include_value=True, fill_cache=True):
        return RocksDBIterator(
            self.db, reverse=reverse, start=start, stop=stop, include_start=include_start, include_stop=include_stop,
            prefix=prefix, include_key=include_key, include_value=include_value
        )

    def write_batch(self, disable_wal: bool = False, sync: bool = False):
        return RocksDBWriteBatch(self.db, sync=sync, disable_wal=disable_wal)

    def close(self):
        self.db.close()
        self.db = None

    @property
    def closed(self) -> bool:
        return self.db is None

    def try_catch_up_with_primary(self):
        self.db.try_catch_up_with_primary()


class RocksDBWriteBatch:
    def __init__(self, db: rocksdb.DB, sync: bool = False, disable_wal: bool = False):
        self.batch = rocksdb.WriteBatch()
        self.db = db
        self.sync = sync
        self.disable_wal = disable_wal

    def __enter__(self):
        return self.batch

    def __exit__(self, exc_type, exc_val, exc_tb):
        if not exc_val:
            self.db.write(self.batch, sync=self.sync, disable_wal=self.disable_wal)


class RocksDBIterator:
    """An iterator for RocksDB."""

    __slots__ = [
        'start',
        'prefix',
        'stop',
        'iterator',
        'include_key',
        'include_value',
        'prev_k',
        'reverse',
        'include_start',
        'include_stop'
    ]

    def __init__(self, db: rocksdb.DB, prefix: bytes = None, start: bool = None, stop: bytes = None,
                 include_key: bool = True, include_value: bool = True, reverse: bool = False,
                 include_start: bool = True, include_stop: bool = False):
        assert (start is None and stop is None) or (prefix is None), 'cannot use start/stop and prefix'
        self.start = start
        self.prefix = prefix
        self.stop = stop
        self.iterator = db.iteritems() if not reverse else reversed(db.iteritems())
        if prefix is not None:
            self.iterator.seek(prefix)
        elif start is not None:
            self.iterator.seek(start)
        self.include_key = include_key
        self.include_value = include_value
        self.prev_k = None
        self.reverse = reverse
        self.include_start = include_start
        self.include_stop = include_stop

    def __iter__(self):
        return self

    def _check_stop_iteration(self, key: bytes):
        if self.stop is not None and (key.startswith(self.stop) or self.stop < key[:len(self.stop)]):
            raise StopIteration
        elif self.start is not None and self.start > key[:len(self.start)]:
            raise StopIteration
        elif self.prefix is not None and not key.startswith(self.prefix):
            raise StopIteration

    def __next__(self):
        # TODO: include start/stop on/off
        # check for needing to stop from previous iteration
        if self.prev_k is not None:
            self._check_stop_iteration(self.prev_k)
        k, v = next(self.iterator)
        self._check_stop_iteration(k)
        self.prev_k = k

        if self.include_key and self.include_value:
            return k, v
        elif self.include_key:
            return k
        return v


class PrefixDB:
    UNDO_KEY_STRUCT = struct.Struct(b'>Q')

    def __init__(self, db: RocksDBStore, max_undo_depth: int = 200, unsafe_prefixes=None):
        self._db = db
        self._op_stack = RevertableOpStack(db.get, unsafe_prefixes=unsafe_prefixes)
        self._max_undo_depth = max_undo_depth

    def unsafe_commit(self):
        """
        Write staged changes to the database without keeping undo information
        Changes written cannot be undone
        """
        try:
            with self._db.write_batch(sync=True) as batch:
                batch_put = batch.put
                batch_delete = batch.delete
                for staged_change in self._op_stack:
                    if staged_change.is_put:
                        batch_put(staged_change.key, staged_change.value)
                    else:
                        batch_delete(staged_change.key)
        finally:
            self._op_stack.clear()

    def commit(self, height: int):
        """
        Write changes for a block height to the database and keep undo information so that the changes can be reverted
        """
        undo_ops = self._op_stack.get_undo_ops()
        delete_undos = []
        if height > self._max_undo_depth:
            delete_undos.extend(self._db.iterator(
                start=DB_PREFIXES.undo.value + self.UNDO_KEY_STRUCT.pack(0),
                stop=DB_PREFIXES.undo.value + self.UNDO_KEY_STRUCT.pack(height - self._max_undo_depth),
                include_value=False
            ))
        try:
            with self._db.write_batch(sync=True) as batch:
                batch_put = batch.put
                batch_delete = batch.delete
                for staged_change in self._op_stack:
                    if staged_change.is_put:
                        batch_put(staged_change.key, staged_change.value)
                    else:
                        batch_delete(staged_change.key)
                for undo_to_delete in delete_undos:
                    batch_delete(undo_to_delete)
                batch_put(DB_PREFIXES.undo.value + self.UNDO_KEY_STRUCT.pack(height), undo_ops)
        finally:
            self._op_stack.clear()

    def rollback(self, height: int):
        """
        Revert changes for a block height
        """
        undo_key = DB_PREFIXES.undo.value + self.UNDO_KEY_STRUCT.pack(height)
        self._op_stack.apply_packed_undo_ops(self._db.get(undo_key))
        try:
            with self._db.write_batch(sync=True) as batch:
                batch_put = batch.put
                batch_delete = batch.delete
                for staged_change in self._op_stack:
                    if staged_change.is_put:
                        batch_put(staged_change.key, staged_change.value)
                    else:
                        batch_delete(staged_change.key)
                batch_delete(undo_key)
        finally:
            self._op_stack.clear()

    def get(self, key: bytes, fill_cache: bool = True) -> Optional[bytes]:
        return self._db.get(key, fill_cache=fill_cache)

    def iterator(self, reverse=False, start=None, stop=None, include_start=True, include_stop=False, prefix=None,
                 include_key=True, include_value=True, fill_cache=True):
        return self._db.iterator(
            reverse=reverse, start=start, stop=stop, include_start=include_start, include_stop=include_stop,
            prefix=prefix, include_key=include_key, include_value=include_value, fill_cache=fill_cache
        )

    def close(self):
        if not self._db.closed:
            self._db.close()

    def try_catch_up_with_primary(self):
        self._db.try_catch_up_with_primary()

    @property
    def closed(self):
        return self._db.closed

    def stage_raw_put(self, key: bytes, value: bytes):
        self._op_stack.append_op(RevertablePut(key, value))

    def stage_raw_delete(self, key: bytes, value: bytes):
        self._op_stack.append_op(RevertableDelete(key, value))