Clear corrupted RocksDB (#888)

Co-authored-by: Daniil Gusev <daniil@quix.io>

Author: gwaramadze

quixstreams/state/rocksdb/exceptions.py

@@ -4,3 +4,6 @@
 
 
 class ColumnFamilyAlreadyExists(StateError): ...
+
+
+class RocksDBCorruptedError(StateError): ...

quixstreams/state/rocksdb/options.py

@@ -32,6 +32,12 @@ class RocksDBOptions(RocksDBOptionsType):
     :param open_max_retries: number of times to retry opening the database
             if it's locked by another process. To disable retrying, pass 0
     :param open_retry_backoff: number of seconds to wait between each retry.
+    :param on_corrupted_recreate: when True, the corrupted DB will be destroyed
+        if the `use_changelog_topics=True` is also set on the Application.
+        If this option is True, but `use_changelog_topics=False`,
+        the DB won't be destroyed.
+        Note that the application doesn't validate the contents of the changelog topics.
+        Default - `False`.
 
     Please see `rocksdict.Options` for a complete description of other options.
     """
@@ -51,6 +57,7 @@ class RocksDBOptions(RocksDBOptionsType):
     open_max_retries: int = 10
     open_retry_backoff: float = 3.0
     use_fsync: bool = True
+    on_corrupted_recreate: bool = False
 
     def to_options(self) -> rocksdict.Options:
         """

quixstreams/state/rocksdb/partition.py

@@ -22,7 +22,7 @@
 from quixstreams.state.recovery import ChangelogProducer
 from quixstreams.state.serialization import int_from_bytes, int_to_bytes
 
-from .exceptions import ColumnFamilyAlreadyExists
+from .exceptions import ColumnFamilyAlreadyExists, RocksDBCorruptedError
 from .metadata import (
     CHANGELOG_OFFSET_KEY,
 )
@@ -350,11 +350,38 @@ def _open_rocksdict(self) -> Rdict:
         options = self._rocksdb_options
         options.create_if_missing(True)
         options.create_missing_column_families(True)
-        rdict = Rdict(
+        create_rdict = lambda: Rdict(
             path=self._path,
             options=options,
             access_type=AccessType.read_write(),
         )
+        # TODO: Add docs
+
+        try:
+            rdict = create_rdict()
+        except Exception as exc:
+            if not str(exc).startswith("Corruption"):
+                raise
+            elif not self._changelog_producer:
+                raise RocksDBCorruptedError(
+                    f'State store at "{self._path}" is corrupted '
+                    f"and cannot be recovered from the changelog topic: "
+                    "`use_changelog_topics` is set to False."
+                ) from exc
+            elif not self._options.on_corrupted_recreate:
+                raise RocksDBCorruptedError(
+                    f'State store at "{self._path}" is corrupted '
+                    f"but may be recovered from the changelog topic. "
+                    "Pass `rocksdb_options=RocksDBOptions(..., on_corrupted_recreate=True)` "
+                    "to the Application to destroy the corrupted state "
+                    "and recover it from the changelog."
+                ) from exc
+
+            logger.warning(f"Destroying corrupted RocksDB path={self._path}")
+            Rdict.destroy(self._path)
+            logger.warning(f"Recreating corrupted RocksDB path={self._path}")
+            rdict = create_rdict()
+
         # Ensure metadata column family is created without defining it upfront
         try:
             rdict.get_column_family(METADATA_CF_NAME)

quixstreams/state/rocksdb/types.py

@@ -23,5 +23,6 @@ class RocksDBOptionsType(Protocol):
     open_max_retries: int
     open_retry_backoff: float
     use_fsync: bool
+    on_corrupted_recreate: bool
 
     def to_options(self) -> rocksdict.Options: ...

tests/test_quixstreams/test_state/fixtures.py

@@ -134,15 +134,15 @@ def factory(
 
 def rocksdb_partition_factory(tmp_path, changelog_producer_mock):
     def factory(
-        name: str = "db",
+        name: str = "",
         options: Optional[RocksDBOptions] = None,
-        changelog_producer: Optional[ChangelogProducer] = None,
+        changelog_producer: Optional[ChangelogProducer] = changelog_producer_mock,
     ) -> RocksDBStorePartition:
         path = (tmp_path / name).as_posix()
         _options = options or RocksDBOptions(open_max_retries=0, open_retry_backoff=3.0)
         return RocksDBStorePartition(
             path,
-            changelog_producer=changelog_producer or changelog_producer_mock,
+            changelog_producer=changelog_producer,
             options=_options,
         )
 

tests/test_quixstreams/test_state/test_rocksdb/test_rocksdb_partition.py

@@ -11,6 +11,7 @@
     RocksDBOptions,
     RocksDBStorePartition,
 )
+from quixstreams.state.rocksdb.exceptions import RocksDBCorruptedError
 from quixstreams.state.rocksdb.windowed.serialization import append_integer
 
 
@@ -67,6 +68,51 @@ def test_open_arbitrary_exception_fails(self, store_partition_factory):
 
         assert str(raised.value) == "some exception"
 
+    def test_db_corrupted_fails_with_no_changelog(
+        self, store_partition_factory, tmp_path
+    ):
+        # Initialize and corrupt the database by messing with the MANIFEST
+        path = tmp_path.as_posix()
+        Rdict(path=path)
+        next(tmp_path.glob("MANIFEST*")).write_bytes(b"")
+
+        with pytest.raises(
+            RocksDBCorruptedError,
+            match=f'State store at "{path}" is corrupted and cannot be recovered '
+            f"from the changelog topic",
+        ):
+            store_partition_factory(changelog_producer=None)
+
+    def test_db_corrupted_fails_with_on_corrupted_recreate_false(
+        self, store_partition_factory, tmp_path
+    ):
+        # Initialize and corrupt the database by messing with the MANIFEST
+        path = tmp_path.as_posix()
+        Rdict(path=path)
+        next(tmp_path.glob("MANIFEST*")).write_bytes(b"")
+
+        with pytest.raises(
+            RocksDBCorruptedError,
+            match=f'State store at "{path}" is corrupted but may be recovered '
+            f"from the changelog topic",
+        ):
+            store_partition_factory()
+
+    def test_db_corrupted_manifest_file(self, store_partition_factory, tmp_path):
+        Rdict(path=tmp_path.as_posix())  # initialize db
+        next(tmp_path.glob("MANIFEST*")).write_bytes(b"")  # write random bytes
+
+        store_partition_factory(options=RocksDBOptions(on_corrupted_recreate=True))
+
+    def test_db_corrupted_sst_file(self, store_partition_factory, tmp_path):
+        rdict = Rdict(path=tmp_path.as_posix())  # initialize db
+        rdict[b"key"] = b"value"  # write something
+        rdict.flush()  # flush creates .sst file
+        rdict.close()  # required to release the lock
+        next(tmp_path.glob("*.sst")).unlink()  # delete the .sst file
+
+        store_partition_factory(options=RocksDBOptions(on_corrupted_recreate=True))
+
     def test_create_and_get_column_family(self, store_partition: RocksDBStorePartition):
         store_partition.create_column_family("cf")
         assert store_partition.get_column_family("cf")
@@ -123,7 +169,7 @@ def test_custom_options(self, store_partition_factory, tmp_path):
         Pass custom "logs_dir" to Rdict and ensure it exists and has some files
         """
 
-        logs_dir = Path(tmp_path / "db" / "logs")
+        logs_dir = Path(tmp_path / "logs")
         options = RocksDBOptions(db_log_dir=logs_dir.as_posix())
         with store_partition_factory(options=options):
             assert logs_dir.is_dir()