220 lines
8 KiB
Python
220 lines
8 KiB
Python
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))
|