import struct import typing import rocksdb from typing import Optional from scribe.db.common import DB_PREFIXES, COLUMN_SETTINGS from scribe.db.revertable import RevertableOpStack, RevertablePut, RevertableDelete ROW_TYPES = {} class PrefixRowType(type): def __new__(cls, name, bases, kwargs): klass = super().__new__(cls, name, bases, kwargs) if name != "PrefixRow": ROW_TYPES[klass.prefix] = klass cache_size = klass.cache_size COLUMN_SETTINGS[klass.prefix] = { 'cache_size': cache_size, } return klass class PrefixRow(metaclass=PrefixRowType): prefix: bytes key_struct: struct.Struct value_struct: struct.Struct key_part_lambdas = [] cache_size: int = 1024 * 1024 * 16 def __init__(self, db: 'rocksdb.DB', op_stack: RevertableOpStack): self._db = db self._op_stack = op_stack self._column_family = self._db.get_column_family(self.prefix) if not self._column_family.is_valid: raise RuntimeError('column family is not valid') def iterate(self, prefix=None, start=None, stop=None, reverse: bool = False, include_key: bool = True, include_value: bool = True, fill_cache: bool = True, deserialize_key: bool = True, deserialize_value: bool = True): if not prefix and not start and not stop: prefix = () if prefix is not None: prefix = self.pack_partial_key(*prefix) if stop is None: try: stop = (int.from_bytes(prefix, byteorder='big') + 1).to_bytes(len(prefix), byteorder='big') except OverflowError: stop = (int.from_bytes(prefix, byteorder='big') + 1).to_bytes(len(prefix) + 1, byteorder='big') else: stop = self.pack_partial_key(*stop) else: if start is not None: start = self.pack_partial_key(*start) if stop is not None: stop = self.pack_partial_key(*stop) if deserialize_key: key_getter = lambda _k: self.unpack_key(_k) else: key_getter = lambda _k: _k if deserialize_value: value_getter = lambda _v: self.unpack_value(_v) else: value_getter = lambda _v: _v it = self._db.iterator( start or prefix, self._column_family, iterate_lower_bound=(start or prefix), iterate_upper_bound=stop, reverse=reverse, include_key=include_key, include_value=include_value, fill_cache=fill_cache, prefix_same_as_start=False ) if include_key and include_value: for k, v in it: yield key_getter(k[1]), value_getter(v) elif include_key: for k in it: yield key_getter(k[1]) elif include_value: for v in it: yield value_getter(v) else: for _ in it: yield None def get(self, *key_args, fill_cache=True, deserialize_value=True): v = self._db.get((self._column_family, self.pack_key(*key_args)), fill_cache=fill_cache) if v: return v if not deserialize_value else self.unpack_value(v) def multi_get(self, key_args: typing.List[typing.Tuple], fill_cache=True, deserialize_value=True): packed_keys = {tuple(args): self.pack_key(*args) for args in key_args} result = { k[-1]: v for k, v in ( self._db.multi_get([(self._column_family, packed_keys[tuple(args)]) for args in key_args], fill_cache=fill_cache) or {} ).items() } def handle_value(v): return None if v is None else v if not deserialize_value else self.unpack_value(v) return [ handle_value(result[packed_keys[tuple(k_args)]]) for k_args in key_args ] def get_pending(self, *key_args, fill_cache=True, deserialize_value=True): packed_key = self.pack_key(*key_args) last_op = self._op_stack.get_last_op_for_key(packed_key) if last_op: if last_op.is_put: return last_op.value if not deserialize_value else self.unpack_value(last_op.value) else: # it's a delete return v = self._db.get((self._column_family, packed_key), fill_cache=fill_cache) if v: return v if not deserialize_value else self.unpack_value(v) def stage_put(self, key_args=(), value_args=()): self._op_stack.append_op(RevertablePut(self.pack_key(*key_args), self.pack_value(*value_args))) def stage_delete(self, key_args=(), value_args=()): self._op_stack.append_op(RevertableDelete(self.pack_key(*key_args), self.pack_value(*value_args))) @classmethod def pack_partial_key(cls, *args) -> bytes: return cls.prefix + cls.key_part_lambdas[len(args)](*args) @classmethod def pack_key(cls, *args) -> bytes: return cls.prefix + cls.key_struct.pack(*args) @classmethod def pack_value(cls, *args) -> bytes: return cls.value_struct.pack(*args) @classmethod def unpack_key(cls, key: bytes): assert key[:1] == cls.prefix, f"prefix should be {cls.prefix}, got {key[:1]}" return cls.key_struct.unpack(key[1:]) @classmethod def unpack_value(cls, data: bytes): return cls.value_struct.unpack(data) @classmethod def unpack_item(cls, key: bytes, value: bytes): return cls.unpack_key(key), cls.unpack_value(value) def estimate_num_keys(self) -> int: return int(self._db.get_property(b'rocksdb.estimate-num-keys', self._column_family).decode()) class BasePrefixDB: """ Base class for a revertable rocksdb database (a rocksdb db where each set of applied changes can be undone) """ UNDO_KEY_STRUCT = struct.Struct(b'>Q32s') PARTIAL_UNDO_KEY_STRUCT = struct.Struct(b'>Q') def __init__(self, path, max_open_files=64, secondary_path='', max_undo_depth: int = 200, unsafe_prefixes=None): column_family_options = {} for prefix in DB_PREFIXES: settings = COLUMN_SETTINGS[prefix.value] column_family_options[prefix.value] = rocksdb.ColumnFamilyOptions() column_family_options[prefix.value].table_factory = rocksdb.BlockBasedTableFactory( block_cache=rocksdb.LRUCache(settings['cache_size']), ) self.column_families: typing.Dict[bytes, 'rocksdb.ColumnFamilyHandle'] = {} options = rocksdb.Options( create_if_missing=True, use_fsync=False, target_file_size_base=33554432, max_open_files=max_open_files if not secondary_path else -1, create_missing_column_families=True ) self._db = rocksdb.DB( path, options, secondary_name=secondary_path, column_families=column_family_options ) for prefix in DB_PREFIXES: cf = self._db.get_column_family(prefix.value) if cf is None and not secondary_path: self._db.create_column_family(prefix.value, column_family_options[prefix.value]) cf = self._db.get_column_family(prefix.value) self.column_families[prefix.value] = cf self._op_stack = RevertableOpStack(self.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: if not len(self._op_stack): return with self._db.write_batch(sync=True) as batch: batch_put = batch.put batch_delete = batch.delete get_column_family = self.column_families.__getitem__ for staged_change in self._op_stack: column_family = get_column_family(DB_PREFIXES(staged_change.key[:1]).value) if staged_change.is_put: batch_put((column_family, staged_change.key), staged_change.value) else: batch_delete((column_family, staged_change.key)) finally: self._op_stack.clear() def commit(self, height: int, block_hash: bytes): """ 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.PARTIAL_UNDO_KEY_STRUCT.pack(0), iterate_upper_bound=DB_PREFIXES.undo.value + self.PARTIAL_UNDO_KEY_STRUCT.pack(height - self._max_undo_depth), include_value=False )) try: undo_c_f = self.column_families[DB_PREFIXES.undo.value] with self._db.write_batch(sync=True) as batch: batch_put = batch.put batch_delete = batch.delete get_column_family = self.column_families.__getitem__ for staged_change in self._op_stack: column_family = get_column_family(DB_PREFIXES(staged_change.key[:1]).value) if staged_change.is_put: batch_put((column_family, staged_change.key), staged_change.value) else: batch_delete((column_family, staged_change.key)) for undo_to_delete in delete_undos: batch_delete((undo_c_f, undo_to_delete)) batch_put((undo_c_f, DB_PREFIXES.undo.value + self.UNDO_KEY_STRUCT.pack(height, block_hash)), undo_ops) finally: self._op_stack.clear() def rollback(self, height: int, block_hash: bytes): """ Revert changes for a block height """ undo_key = DB_PREFIXES.undo.value + self.UNDO_KEY_STRUCT.pack(height, block_hash) undo_c_f = self.column_families[DB_PREFIXES.undo.value] undo_info = self._db.get((undo_c_f, undo_key)) self._op_stack.apply_packed_undo_ops(undo_info) try: with self._db.write_batch(sync=True) as batch: batch_put = batch.put batch_delete = batch.delete get_column_family = self.column_families.__getitem__ for staged_change in self._op_stack: column_family = get_column_family(DB_PREFIXES(staged_change.key[:1]).value) if staged_change.is_put: batch_put((column_family, staged_change.key), staged_change.value) else: batch_delete((column_family, staged_change.key)) # batch_delete(undo_key) finally: self._op_stack.clear() def get(self, key: bytes, fill_cache: bool = True) -> Optional[bytes]: cf = self.column_families[key[:1]] return self._db.get((cf, key), fill_cache=fill_cache) def iterator(self, start: bytes, column_family: 'rocksdb.ColumnFamilyHandle' = None, iterate_lower_bound: bytes = None, iterate_upper_bound: bytes = None, reverse: bool = False, include_key: bool = True, include_value: bool = True, fill_cache: bool = True, prefix_same_as_start: bool = False, auto_prefix_mode: bool = True): return self._db.iterator( start=start, column_family=column_family, iterate_lower_bound=iterate_lower_bound, iterate_upper_bound=iterate_upper_bound, reverse=reverse, include_key=include_key, include_value=include_value, fill_cache=fill_cache, prefix_same_as_start=prefix_same_as_start, auto_prefix_mode=auto_prefix_mode ) def close(self): self._db.close() def try_catch_up_with_primary(self): self._db.try_catch_up_with_primary() 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)) def estimate_num_keys(self, column_family: 'rocksdb.ColumnFamilyHandle' = None): return int(self._db.get_property(b'rocksdb.estimate-num-keys', column_family).decode())