Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add --temp-data CLI option #1228

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions hathor/builder/cli_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,23 +131,28 @@ def create_manager(self, reactor: Reactor) -> HathorManager:
self.event_ws_factory: Optional[EventWebsocketFactory] = None

if self._args.memory_storage:
self.log.warn('--memory-storage is deprecated, use --temp-data instead')
self.check_or_raise(not self._args.data, '--data should not be used with --memory-storage')
self.check_or_raise(not self._args.temp_data, '--temp-data should not be used with --memory-storage')
# if using MemoryStorage, no need to have cache
indexes = MemoryIndexesManager()
tx_storage = TransactionMemoryStorage(indexes, settings=settings)
event_storage = EventMemoryStorage()
self.check_or_raise(not self._args.x_rocksdb_indexes, 'RocksDB indexes require RocksDB data')
self.log.info('with storage', storage_class=type(tx_storage).__name__)
else:
self.check_or_raise(bool(self._args.data), '--data is expected')
assert self._args.data is not None
self.check_or_raise(
bool(self._args.data) or self._args.temp_data,
'either --data or --temp-data is expected'
)
if self._args.rocksdb_storage:
self.log.warn('--rocksdb-storage is now implied, no need to specify it')
cache_capacity = self._args.rocksdb_cache
self.rocksdb_storage = RocksDBStorage(path=self._args.data, cache_capacity=cache_capacity)

# Initialize indexes manager.
if self._args.memory_indexes:
self.log.warn('--memory-indexes is deprecated')
indexes = MemoryIndexesManager()
else:
indexes = RocksDBIndexesManager(self.rocksdb_storage)
Expand Down
10 changes: 6 additions & 4 deletions hathor/cli/run_node.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,12 +95,14 @@ def create_parser(cls) -> ArgumentParser:
parser.add_argument('--x-status-ipv6-interface', help='IPv6 interface to bind the status server')
parser.add_argument('--stratum', type=int, help='Port to run stratum server')
parser.add_argument('--x-stratum-ipv6-interface', help='IPv6 interface to bind the stratum server')
parser.add_argument('--data', help='Data directory')
data_group = parser.add_mutually_exclusive_group()
data_group.add_argument('--data', help='Data directory')
data_group.add_argument('--temp-data', action='store_true',
help='Automatically create storage in a temporary directory')
storage = parser.add_mutually_exclusive_group()
storage.add_argument('--rocksdb-storage', action='store_true', help='Use RocksDB storage backend (default)')
storage.add_argument('--memory-storage', action='store_true', help='Do not use a persistent storage')
parser.add_argument('--memory-indexes', action='store_true',
help='Use memory indexes when using RocksDB storage (startup is significantly slower)')
storage.add_argument('--memory-storage', action='store_true', help=SUPPRESS) # deprecated
parser.add_argument('--memory-indexes', action='store_true', help=SUPPRESS) # deprecated
parser.add_argument('--rocksdb-cache', type=int, help='RocksDB block-table cache size (bytes)', default=None)
parser.add_argument('--wallet', help='Set wallet type. Options are hd (Hierarchical Deterministic) or keypair',
default=None)
Expand Down
1 change: 1 addition & 0 deletions hathor/cli/run_node_args.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ class RunNodeArgs(BaseModel, extra=Extra.allow):
rocksdb_storage: bool
memory_storage: bool
memory_indexes: bool
temp_data: bool
rocksdb_cache: Optional[int]
wallet: Optional[str]
wallet_enable_api: bool
Expand Down
40 changes: 30 additions & 10 deletions hathor/storage/rocksdb_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from typing import TYPE_CHECKING, Optional
from __future__ import annotations

if TYPE_CHECKING:
import rocksdb
import os
import tempfile

import rocksdb
from structlog import get_logger
from typing_extensions import assert_never

logger = get_logger()
_DB_NAME = 'data_v2.db'
Expand All @@ -28,12 +29,16 @@ class RocksDBStorage:
""" Creates a RocksDB database
Give clients the option to create column families
"""
def __init__(self, path: str = './', cache_capacity: Optional[int] = None):
import rocksdb
def __init__(
self,
path: str | tempfile.TemporaryDirectory | None = None,
cache_capacity: int | None = None,
) -> None:
self.log = logger.new()
self._path = path
# We have to keep a reference to the TemporaryDirectory because it is cleaned up when garbage collected.
self.path, self.temp_dir = self._get_path_and_temp_dir(path)

db_path = os.path.join(path, _DB_NAME)
db_path = os.path.join(self.path, _DB_NAME)
lru_cache = cache_capacity and rocksdb.LRUCache(cache_capacity)
table_factory = rocksdb.BlockBasedTableFactory(block_cache=lru_cache)
options = rocksdb.Options(
Expand All @@ -58,13 +63,28 @@ def __init__(self, path: str = './', cache_capacity: Optional[int] = None):

# finally, open the database
self._db = rocksdb.DB(db_path, options, column_families=column_families)
self.log.info('starting rocksdb', path=self.path)
self.log.debug('open db', cf_list=[cf.name.decode('ascii') for cf in self._db.column_families])

def get_db(self) -> 'rocksdb.DB':
@staticmethod
def _get_path_and_temp_dir(
path: str | tempfile.TemporaryDirectory | None,
) -> tuple[str, tempfile.TemporaryDirectory | None]:
match path:
case str():
return path, None
case tempfile.TemporaryDirectory():
return path.name, path
case None:
temp_dir = tempfile.TemporaryDirectory()
return temp_dir.name, temp_dir
case _:
assert_never(path)

def get_db(self) -> rocksdb.DB:
return self._db

def get_or_create_column_family(self, cf_name: bytes) -> 'rocksdb.ColumnFamilyHandle':
import rocksdb
cf = self._db.get_column_family(cf_name)
if cf is None:
cf = self._db.create_column_family(cf_name, rocksdb.ColumnFamilyOptions())
Expand Down
2 changes: 1 addition & 1 deletion tests/others/test_cli_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ def _build(self, cmd_args: list[str]) -> HathorManager:
return manager

def test_empty(self):
self._build_with_error([], '--data is expected')
self._build_with_error([], 'either --data or --temp-data is expected')

def test_all_default(self):
data_dir = self.mkdtemp()
Expand Down