Base implementation of RocksDB support (#1416)

* Add RocksDB backend

* Ensure we dont reuse database dir for different database engines

* Cache rocksdb installation

* PR feedback

* fix bash scripts

* fix types

Author: pipermerriam

.circleci/config.yml

@@ -23,29 +23,48 @@ common: &common
         when: on_fail
     - restore_cache:
         keys:
-          - cache-{{ .Environment.CIRCLE_JOB }}-{{ checksum "setup.py" }}-{{ checksum "tox.ini" }}
+          - cache-v1-python-{{ arch }}-{{ .Environment.CIRCLE_JOB }}-{{ checksum "setup.py" }}-{{ checksum "tox.ini" }}
+    - restore_cache:
+        keys:
+          - cache-v1-rocksdb-{{ arch }}-{{ .Environment.CIRCLE_JOB }}-{{ checksum ".circleci/install_rocksdb.sh" }}
+    - run:
+        name: install rocksdb
+        command: sudo sh ./.circleci/install_rocksdb.sh
+    - save_cache:
+        paths:
+          - ~/rocksdb/
+        key: cache-v1-rocksdb-{{ arch }}-{{ .Environment.CIRCLE_JOB }}-{{ checksum ".circleci/install_rocksdb.sh" }}
     - run:
         name: install dependencies
         command: pip install --user tox
     - run:
         name: run tox
-        command: ~/.local/bin/tox
+        command: ~/.local/bin/tox -r
     - save_cache:
         paths:
           - .hypothesis
-          - .tox
           - ~/.cache/pip
           - ~/.local
           - ./eggs
-        key: cache-{{ .Environment.CIRCLE_JOB }}-{{ checksum "setup.py" }}-{{ checksum "tox.ini" }}
+        key: cache-v1-python-{{ arch }}-{{ .Environment.CIRCLE_JOB }}-{{ checksum "setup.py" }}-{{ checksum "tox.ini" }}
 
 geth_steps: &geth_steps
   working_directory: ~/repo
   steps:
     - checkout
     - restore_cache:
         keys:
-          - cache-{{ .Environment.CIRCLE_JOB }}-{{ checksum "setup.py" }}-{{ checksum "tox.ini" }}
+          - cache-v1-python-{{ arch }}-{{ .Environment.CIRCLE_JOB }}-{{ checksum "setup.py" }}-{{ checksum "tox.ini" }}
+    - restore_cache:
+        keys:
+          - cache-v2-rocksdb-{{ arch }}-{{ .Environment.CIRCLE_JOB }}-{{ checksum ".circleci/install_rocksdb.sh" }}
+    - run:
+        name: install rocksdb
+        command: sudo sh ./.circleci/install_rocksdb.sh
+    - save_cache:
+        paths:
+          - ~/rocksdb/
+        key: cache-v2-rocksdb-{{ arch }}-{{ .Environment.CIRCLE_JOB }}-{{ checksum ".circleci/install_rocksdb.sh" }}
     - run:
         name: install dependencies
         command: pip install --user tox
@@ -79,7 +98,7 @@ geth_steps: &geth_steps
           - ./eggs
           - ~/.ethash
           - ~/.py-geth
-        key: cache-{{ .Environment.CIRCLE_JOB }}-{{ checksum "setup.py" }}-{{ checksum "tox.ini" }}
+        key: cache-v1-python-{{ arch }}-{{ .Environment.CIRCLE_JOB }}-{{ checksum "setup.py" }}-{{ checksum "tox.ini" }}
 
 jobs:
   py35-lint:
@@ -266,7 +285,6 @@ jobs:
         environment:
           TOXENV: py36-trinity-integration
   py36-trinity-lightchain_integration:
-    <<: *common
     <<: *geth_steps
     docker:
       - image: circleci/python:3.6

.circleci/install_rocksdb.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+
+set -o errexit
+set -o nounset
+
+sudo apt-get install -y liblz4-dev libsnappy-dev libgflags-dev zlib1g-dev libbz2-dev libzstd-dev
+
+
+if [ ! -d "/home/circleci/rocksdb" ]; then
+  git clone https://github.com/facebook/rocksdb /home/circleci/rocksdb
+fi
+if [ ! -f "/home/circleci/rocksdb/librocksdb.so.5.8.8" ]; then
+  cd /home/circleci/rocksdb/ && git checkout v5.8.8 && sudo make install-shared INSTALL_PATH=/usr
+fi
+if [ ! -f "/usr/lib/librocksdb.so.5.8" ]; then
+  ln -fs /home/circleci/rocksdb/librocksdb.so.5.8.8 /usr/lib/librocksdb.so.5.8
+fi
+if [ ! -f "/usr/lib/librocksdb.so.5" ]; then
+  ln -fs /home/circleci/rocksdb/librocksdb.so.5.8.8 /usr/lib/librocksdb.so.5
+fi
+if [ ! -f "/usr/lib/librocksdb.so" ]; then
+  ln -fs /home/circleci/rocksdb/librocksdb.so.5.8.8 /usr/lib/librocksdb.so
+fi  

.circleci/merge_pr.sh

@@ -1,5 +1,8 @@
 #!/usr/bin/env bash
 
+set -o errexit
+set -o nounset
+
 if [[ -n "${CIRCLE_PR_NUMBER}" ]]; then
   PR_INFO_URL=https://api.github.com/repos/$CIRCLE_PROJECT_USERNAME/$CIRCLE_PROJECT_REPONAME/pulls/$CIRCLE_PR_NUMBER
   PR_BASE_BRANCH=$(curl -L "$PR_INFO_URL" | python -c 'import json, sys; obj = json.load(sys.stdin); sys.stdout.write(obj["base"]["ref"])')

docs/guides/trinity/quickstart.rst

@@ -26,6 +26,12 @@ we need to install the ``python3-pip`` package through the following command.
 
   apt-get install python3-pip
 
+Trinity also requires RocksDB which can be installed with the following command:
+
+.. code:: sh
+
+  apt-get install liblz4-dev lib-rocksdb5.8
+
 .. note::
   .. include:: /fragments/virtualenv_explainer.rst
 
@@ -42,7 +48,7 @@ First, install LevelDB and the latest Python 3 with brew:
 
 .. code:: sh
 
-  brew install python3 leveldb
+  brew install python3 leveldb rocksdb
 
 .. note::
   .. include:: /fragments/virtualenv_explainer.rst

eth/db/backends/level.py

@@ -2,7 +2,7 @@
 import logging
 from pathlib import Path
 from typing import (
-    Generator,
+    Iterator,
     TYPE_CHECKING,
 )
 
@@ -27,10 +27,9 @@
 class LevelDB(BaseAtomicDB):
     logger = logging.getLogger("eth.db.backends.LevelDB")
 
-    # Creates db as a class variable to avoid level db lock error
     def __init__(self, db_path: Path = None) -> None:
         if not db_path:
-            raise TypeError("Please specifiy a valid path for your database.")
+            raise TypeError("The LevelDB backend requires a database path")
         try:
             with catch_and_ignore_import_warning():
                 import plyvel  # noqa: F811
@@ -54,10 +53,13 @@ def _exists(self, key: bytes) -> bool:
         return self.db.get(key) is not None
 
     def __delitem__(self, key: bytes) -> None:
+        v = self.db.get(key)
+        if v is None:
+            raise KeyError(key)
         self.db.delete(key)
 
     @contextmanager
-    def atomic_batch(self) -> Generator['LevelDBWriteBatch', None, None]:
+    def atomic_batch(self) -> Iterator['LevelDBWriteBatch']:
         with self.db.write_batch(transaction=True) as atomic_batch:
             readable_batch = LevelDBWriteBatch(self, atomic_batch)
             try:

eth/db/backends/rocks.py

@@ -0,0 +1,139 @@
+from contextlib import contextmanager
+import logging
+from pathlib import Path
+from typing import (
+    Iterator,
+    TYPE_CHECKING,
+)
+
+from eth_utils import ValidationError
+
+from eth.db.diff import (
+    DBDiffTracker,
+    DiffMissingError,
+)
+from .base import (
+    BaseAtomicDB,
+    BaseDB,
+)
+
+if TYPE_CHECKING:
+    import rocksdb  # noqa: F401
+
+
+class RocksDB(BaseAtomicDB):
+    logger = logging.getLogger("eth.db.backends.RocksDB")
+
+    def __init__(self,
+                 db_path: Path = None,
+                 opts: 'rocksdb.Options' = None,
+                 read_only: bool=False) -> None:
+        if not db_path:
+            raise TypeError("The RocksDB backend requires a database path")
+        try:
+            import rocksdb  # noqa: F811
+        except ImportError:
+            raise ImportError(
+                "RocksDB requires the python-rocksdb library which is not "
+                "available for import."
+            )
+
+        if opts is None:
+            opts = rocksdb.Options(create_if_missing=True)
+        self.db_path = db_path
+        self.db = rocksdb.DB(str(db_path), opts, read_only=read_only)
+
+    def __getitem__(self, key: bytes) -> bytes:
+        v = self.db.get(key)
+        if v is None:
+            raise KeyError(key)
+        return v
+
+    def __setitem__(self, key: bytes, value: bytes) -> None:
+        self.db.put(key, value)
+
+    def _exists(self, key: bytes) -> bool:
+        return self.db.get(key) is not None
+
+    def __delitem__(self, key: bytes) -> None:
+        exists, _ = self.db.key_may_exist(key)
+        if not exists:
+            raise KeyError(key)
+        self.db.delete(key)
+
+    @contextmanager
+    def atomic_batch(self) -> Iterator['RocksDBWriteBatch']:
+        import rocksdb  # noqa: F811
+        batch = rocksdb.WriteBatch()
+
+        readable_batch = RocksDBWriteBatch(self, batch)
+
+        try:
+            yield readable_batch
+        finally:
+            readable_batch.decommission()
+
+        self.db.write(batch)
+
+
+class RocksDBWriteBatch(BaseDB):
+    """
+    A native rocksdb write batch does not permit reads on the in-progress data.
+    This class fills that gap, by tracking the in-progress diff, and adding
+    a read interface.
+    """
+    logger = logging.getLogger("eth.db.backends.RocksDBWriteBatch")
+
+    def __init__(self, original_read_db: BaseDB, write_batch: 'rocksdb.WriteBatch') -> None:
+        self._original_read_db = original_read_db
+        self._write_batch = write_batch
+        # keep track of the temporary changes made
+        self._track_diff = DBDiffTracker()
+
+    def __getitem__(self, key: bytes) -> bytes:
+        if self._track_diff is None:
+            raise ValidationError("Cannot get data from a write batch, out of context")
+
+        try:
+            changed_value = self._track_diff[key]
+        except DiffMissingError as missing:
+            if missing.is_deleted:
+                raise KeyError(key)
+            else:
+                return self._original_read_db[key]
+        else:
+            return changed_value
+
+    def __setitem__(self, key: bytes, value: bytes) -> None:
+        if self._track_diff is None:
+            raise ValidationError("Cannot set data from a write batch, out of context")
+
+        self._write_batch.put(key, value)
+        self._track_diff[key] = value
+
+    def _exists(self, key: bytes) -> bool:
+        if self._track_diff is None:
+            raise ValidationError("Cannot test data existance from a write batch, out of context")
+
+        try:
+            self._track_diff[key]
+        except DiffMissingError as missing:
+            if missing.is_deleted:
+                return False
+            else:
+                return key in self._original_read_db
+        else:
+            return True
+
+    def __delitem__(self, key: bytes) -> None:
+        if self._track_diff is None:
+            raise ValidationError("Cannot delete data from a write batch, out of context")
+
+        self._write_batch.delete(key)
+        del self._track_diff[key]
+
+    def decommission(self) -> None:
+        """
+        Prevent any further actions to be taken on this write batch, called after leaving context
+        """
+        self._track_diff = None

scripts/benchmark/utils/chain_plumbing.py

@@ -29,8 +29,8 @@
 from eth.chains.base import (
     MiningChain,
 )
-from eth.db.backends.level import (
-    LevelDB,
+from eth.db.backends.rocks import (
+    RocksDB,
 )
 from eth.vm.base import (
     BaseVM,
@@ -91,14 +91,14 @@
 def get_chain(vm: Type[BaseVM], genesis_state: GenesisState) -> Iterable[MiningChain]:
 
     with tempfile.TemporaryDirectory() as temp_dir:
-        level_db_obj = LevelDB(Path(temp_dir))
-        level_db_chain = build(
+        base_db = RocksDB(Path(temp_dir))
+        chain = build(
             MiningChain,
             fork_at(vm, constants.GENESIS_BLOCK_NUMBER),
             disable_pow_check(),
-            genesis(db=level_db_obj, params=GENESIS_PARAMS, state=genesis_state)
+            genesis(db=base_db, params=GENESIS_PARAMS, state=genesis_state)
         )
-        yield level_db_chain
+        yield chain
 
 
 def get_all_chains(genesis_state: GenesisState=DEFAULT_GENESIS_STATE) -> Iterable[MiningChain]:

setup.py

@@ -27,6 +27,7 @@
         "eth-hash[pysha3];implementation_name=='cpython'",
         "eth-hash[pycryptodome];implementation_name=='pypy'",
         "plyvel==1.0.5",
+        "python-rocksdb==0.6.9",
     ],
     'p2p': [
         "asyncio-cancel-token==0.1.0a2",
@@ -43,6 +44,7 @@
         "coincurve>=8.0.0,<9.0.0",
         "ipython>=6.2.1,<7.0.0",
         "plyvel==1.0.5",
+        "python-rocksdb==0.6.9",
         "web3==4.4.1",
         "lahja==0.9.0",
         "termcolor>=1.1.0,<2.0.0",
@@ -57,7 +59,7 @@
         "pytest-asyncio==0.9.0",
         "pytest-cov==2.5.1",
         "pytest-watch>=4.1.0,<5",
-        "pytest-xdist==1.18.1",
+        "pytest-xdist==1.23.2",
         # only needed for p2p
         "pytest-asyncio-network-simulator==0.1.0a2;python_version>='3.6'",
     ],

tests/database/test_base_atomic_db.py

@@ -4,14 +4,17 @@
 
 from eth.db.atomic import AtomicDB
 from eth.db.backends.level import LevelDB
+from eth.db.backends.rocks import RocksDB
 
 
-@pytest.fixture(params=['atomic', 'level'])
+@pytest.fixture(params=['atomic', 'level', 'rocks'])
 def atomic_db(request, tmpdir):
     if request.param == 'atomic':
         return AtomicDB()
     elif request.param == 'level':
         return LevelDB(db_path=tmpdir.mkdir("level_db_path"))
+    elif request.param == 'rocks':
+        return RocksDB(db_path=tmpdir.mkdir("rocks_db_path"))
     else:
         raise ValueError("Unexpected database type: {}".format(request.param))