Add unit tests for rocksdict to guarantee reliability (#478)

* initial commit

* create new class for rocksdict testing

* commit progress

* update ci

* fix partition get call

* commit some fixed tests

* commit remaining TODO tests

* commit and finish later

* upgrade rocksdict minimum to 0.3.8

* commit cleanup

* 8 tests to go

* cleanup

* cleanup redundant tests

* cleanup

* cleanup

* undo breaking changes

* commit fixes for more tests

* commit fixes for more tests

* commit remaining tests

* bring test__size back

Author: wbarnha

faust/stores/rocksdb.py

@@ -95,6 +95,7 @@ class RocksDBOptions:
     block_cache_size: int = DEFAULT_BLOCK_CACHE_SIZE
     block_cache_compressed_size: int = DEFAULT_BLOCK_CACHE_COMPRESSED_SIZE
     bloom_filter_size: int = DEFAULT_BLOOM_FILTER_SIZE
+    use_rocksdict: bool = USE_ROCKSDICT
     extra_options: Mapping
 
     def __init__(
@@ -106,6 +107,7 @@ def __init__(
         block_cache_size: Optional[int] = None,
         block_cache_compressed_size: Optional[int] = None,
         bloom_filter_size: Optional[int] = None,
+        use_rocksdict: Optional[bool] = None,
         **kwargs: Any,
     ) -> None:
         if max_open_files is not None:
@@ -122,11 +124,13 @@ def __init__(
             self.block_cache_compressed_size = block_cache_compressed_size
         if bloom_filter_size is not None:
             self.bloom_filter_size = bloom_filter_size
+        if use_rocksdict is not None:
+            self.use_rocksdict = use_rocksdict
         self.extra_options = kwargs
 
     def open(self, path: Path, *, read_only: bool = False) -> DB:
         """Open RocksDB database using this configuration."""
-        if USE_ROCKSDICT:
+        if self.use_rocksdict:
             db_options = self.as_options()
             db_options.set_db_paths(
                 [rocksdict.DBPath(str(path), self.target_file_size_base)]
@@ -139,7 +143,7 @@ def open(self, path: Path, *, read_only: bool = False) -> DB:
 
     def as_options(self) -> Options:
         """Return :class:`rocksdb.Options` object using this configuration."""
-        if USE_ROCKSDICT:
+        if self.use_rocksdict:
             db_options = Options(raw_mode=True)
             db_options.create_if_missing(True)
             db_options.set_max_open_files(self.max_open_files)
@@ -156,9 +160,6 @@ def as_options(self) -> Options:
             table_factory_options.set_index_type(
                 rocksdict.BlockBasedIndexType.binary_search()
             )
-            table_factory_options.set_block_cache_compressed(
-                rocksdict.Cache(self.block_cache_compressed_size)
-            )
             db_options.set_block_based_table_factory(table_factory_options)
             return db_options
         else:
@@ -221,16 +222,13 @@ def __init__(
         driver: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
-        if rocksdict is None and rocksdb is None:
+        if rocksdb is None and rocksdict is None:
             error = ImproperlyConfigured(
                 "RocksDB bindings not installed? pip install faust-streaming-rocksdb"
                 " or rocksdict"
             )
             try:
-                try:
-                    import rocksdb as _rocksdb  # noqa: F401
-                except ImportError:
-                    import rocksdict as _rocksdict  # noqa: F401
+                import rocksdb as _rocksdb  # noqa: F401
             except Exception as exc:  # pragma: no cover
                 raise error from exc
             else:  # pragma: no cover
@@ -242,11 +240,12 @@ def __init__(
         self.read_only = self.options.pop("read_only", read_only)
 
         self.driver = self.options.pop("driver", driver)
-        global USE_ROCKSDICT
         if self.driver == "rocksdict":
-            USE_ROCKSDICT = True
+            self.USE_ROCKSDICT = True
         elif self.driver == "python-rocksdb":
-            USE_ROCKSDICT = False
+            self.USE_ROCKSDICT = False
+        else:
+            self.USE_ROCKSDICT = USE_ROCKSDICT
 
         self.rocksdb_options = RocksDBOptions(**self.options)
         if key_index_size is None:
@@ -353,13 +352,7 @@ def persisted_offset(self, tp: TP) -> Optional[int]:
 
         See :meth:`set_persisted_offset`.
         """
-        if USE_ROCKSDICT:
-            try:
-                offset = self._db_for_partition(tp.partition)[self.offset_key]
-            except Exception:
-                offset = None
-        else:
-            offset = self._db_for_partition(tp.partition).get(self.offset_key)
+        offset = self._db_for_partition(tp.partition).get(self.offset_key)
         if offset is not None:
             return int(offset)
         return None
@@ -372,12 +365,7 @@ def set_persisted_offset(self, tp: TP, offset: int) -> None:
         to only read the events that occurred recently while
         we were not an active replica.
         """
-        if USE_ROCKSDICT:
-            self._db_for_partition(tp.partition)[self.offset_key] = str(offset).encode()
-        else:
-            self._db_for_partition(tp.partition).put(
-                self.offset_key, str(offset).encode()
-            )
+        self._db_for_partition(tp.partition).put(self.offset_key, str(offset).encode())
 
     async def need_active_standby_for(self, tp: TP) -> bool:
         """Decide if an active standby is needed for this topic partition.
@@ -425,7 +413,7 @@ def apply_changelog_batch(
                 of a changelog event.
         """
         batches: DefaultDict[int, WriteBatch]
-        if USE_ROCKSDICT:
+        if self.USE_ROCKSDICT:
             batches = defaultdict(lambda: WriteBatch(raw_mode=True))
         else:
             batches = defaultdict(rocksdb.WriteBatch)
@@ -436,9 +424,13 @@ def apply_changelog_batch(
                 offset if tp not in tp_offsets else max(offset, tp_offsets[tp])
             )
             msg = event.message
+            if self.USE_ROCKSDICT:
+                msg.key = msg.key.encode()
             if msg.value is None:
                 batches[msg.partition].delete(msg.key)
             else:
+                if self.USE_ROCKSDICT:
+                    msg.value = msg.value.encode()
                 batches[msg.partition].put(msg.key, msg.value)
 
         for partition, batch in batches.items():
@@ -453,10 +445,7 @@ def _set(self, key: bytes, value: Optional[bytes]) -> None:
         partition = event.message.partition
         db = self._db_for_partition(partition)
         self._key_index[key] = partition
-        if USE_ROCKSDICT:
-            db[key] = value
-        else:
-            db.put(key, value)
+        db.put(key, value)
 
     def _db_for_partition(self, partition: int) -> DB:
         try:
@@ -481,13 +470,7 @@ def _get(self, key: bytes) -> Optional[bytes]:
         if partition_from_message:
             partition = event.message.partition
             db = self._db_for_partition(partition)
-            if USE_ROCKSDICT:
-                try:
-                    value = db[key]
-                except Exception:
-                    value = None
-            else:
-                value = db.get(key)
+            value = db.get(key)
             if value is not None:
                 self._key_index[key] = partition
             return value
@@ -511,11 +494,9 @@ def _get_bucket_for_key(self, key: bytes) -> Optional[_DBValueTuple]:
             dbs = cast(Iterable[PartitionDB], self._dbs.items())
 
         for partition, db in dbs:
-            if USE_ROCKSDICT:
-                try:
-                    value = db[key]
-                except Exception:
-                    value = None
+            if self.USE_ROCKSDICT:
+                # TODO: Remove this once key_may_exist is added
+                value = db.get(key)
                 if value is not None:
                     self._key_index[key] = partition
                     return _DBValueTuple(db, value)
@@ -640,25 +621,19 @@ def _contains(self, key: bytes) -> bool:
         if partition_from_message:
             partition = event.message.partition
             db = self._db_for_partition(partition)
-            if USE_ROCKSDICT:
-                try:
-                    value = db[key]
-                except Exception:
-                    value = None
-            else:
-                value = db.get(key)
+            value = db.get(key)
             if value is not None:
                 return True
             else:
                 return False
         else:
             for db in self._dbs_for_key(key):
                 # bloom filter: false positives possible, but not false negatives
-                if USE_ROCKSDICT:
-                    try:
-                        _ = db[key]
+                if self.USE_ROCKSDICT:
+                    # TODO: Remove once key_may_exist is added
+                    if db.get(key) is not None:
                         return True
-                    except Exception:
+                    else:
                         return False
                 else:
                     if db.key_may_exist(key)[0] and db.get(key) is not None:
@@ -687,7 +662,7 @@ def _size(self) -> int:
         return sum(self._size1(db) for db in self._dbs_for_actives())
 
     def _visible_keys(self, db: DB) -> Iterator[bytes]:
-        if USE_ROCKSDICT:
+        if self.USE_ROCKSDICT:
             it = db.keys()
             iter = db.iter()
             iter.seek_to_first()
@@ -699,7 +674,7 @@ def _visible_keys(self, db: DB) -> Iterator[bytes]:
                 yield key
 
     def _visible_items(self, db: DB) -> Iterator[Tuple[bytes, bytes]]:
-        if USE_ROCKSDICT:
+        if self.USE_ROCKSDICT:
             it = db.items()
         else:
             it = db.iteritems()  # noqa: B301

requirements/extras/rocksdict.txt

@@ -1 +1 @@
-rocksdict>=0.3.7
+rocksdict>=0.3.8