Support Version Rollbacks for RocksDB (#6)

plugin-api/build.gradle

@@ -56,7 +56,7 @@ Calculated : ${currentHash}
 tasks.register('checkAPIChanges', FileStateChecker) {
   description = "Checks that the API for the Plugin-API project does not change without deliberate thought"
   files = sourceSets.main.allJava.files
-  knownHash = 'mBaqR6fbWYPndHOw7CymPkR3KKb+f8pxLldBnTCjYAg='
+  knownHash = '+dpM9DHxYq73BGIktTSVRKPjZ1mu+Nw8NoYYvPgmSn4='
 }
 check.dependsOn('checkAPIChanges')
 

plugin-api/src/main/java/org/hyperledger/besu/plugin/services/storage/KeyValueStorageFactory.java

@@ -17,9 +17,11 @@
 import org.hyperledger.besu.plugin.services.MetricsSystem;
 import org.hyperledger.besu.plugin.services.exception.StorageException;
 
+import java.io.Closeable;
+
 /** Factory for creating key-value storage instances. */
 @Unstable
-public interface KeyValueStorageFactory {
+public interface KeyValueStorageFactory extends Closeable {
 
   /**
    * Retrieves the identity of the key-value storage factory.

plugins/rocksdb/src/main/java/org/hyperledger/besu/plugin/services/storage/rocksdb/RocksDBKeyValueStorageFactory.java

@@ -12,6 +12,8 @@
  */
 package org.hyperledger.besu.plugin.services.storage.rocksdb;
 
+import static com.google.common.base.Preconditions.checkNotNull;
+
 import org.hyperledger.besu.plugin.services.BesuConfiguration;
 import org.hyperledger.besu.plugin.services.MetricsSystem;
 import org.hyperledger.besu.plugin.services.exception.StorageException;
@@ -40,22 +42,36 @@
 public class RocksDBKeyValueStorageFactory implements KeyValueStorageFactory {
 
   private static final Logger LOG = LogManager.getLogger();
-  private static final int DEFAULT_VERSION = 1;
-  private static final Set<Integer> SUPPORTED_VERSION = Set.of(0, 1);
+  private final int DEFAULT_VERSION;
+  private static final Set<Integer> SUPPORTED_VERSIONS = Set.of(0, 1);
   private static final String NAME = "rocksdb";
 
-  private boolean isSegmentIsolationSupported;
+  private Integer databaseVersion;
+  private Boolean isSegmentIsolationSupported;
   private SegmentedKeyValueStorage<?> segmentedStorage;
   private KeyValueStorage unsegmentedStorage;
+  private RocksDBConfiguration rocksDBConfiguration;
 
   private final Supplier<RocksDBFactoryConfiguration> configuration;
   private final List<SegmentIdentifier> segments;
 
-  public RocksDBKeyValueStorageFactory(
+  RocksDBKeyValueStorageFactory(
       final Supplier<RocksDBFactoryConfiguration> configuration,
-      final List<SegmentIdentifier> segments) {
+      final List<SegmentIdentifier> segments,
+      final int DEFAULT_VERSION) {
     this.configuration = configuration;
     this.segments = segments;
+    this.DEFAULT_VERSION = DEFAULT_VERSION;
+  }
+
+  public RocksDBKeyValueStorageFactory(
+      final Supplier<RocksDBFactoryConfiguration> configuration,
+      final List<SegmentIdentifier> segments) {
+    this(
+        configuration,
+        segments,
+        /** Source of truth for the default database version. */
+        1);
   }
 
   @Override
@@ -71,62 +87,64 @@ public KeyValueStorage create(
       throws StorageException {
 
     if (requiresInit()) {
-      init(commonConfiguration, metricsSystem);
+      init(commonConfiguration);
     }
 
-    return isSegmentIsolationSupported
-        ? new SegmentedKeyValueStorageAdapter<>(segment, segmentedStorage)
-        : unsegmentedStorage;
-  }
-
-  @Override
-  public boolean isSegmentIsolationSupported() {
-    return isSegmentIsolationSupported;
-  }
-
-  public void close() throws IOException {
-    if (segmentedStorage != null) {
-      segmentedStorage.close();
-    }
-    if (unsegmentedStorage != null) {
-      unsegmentedStorage.close();
+    // It's probably a good idea for the creation logic to be entirely dependent on the database
+    // version. Introducing intermediate booleans that represent database properties and dispatching
+    // creation logic based on them is error prone.
+    switch (databaseVersion) {
+      case 0:
+        {
+          segmentedStorage = null;
+          if (unsegmentedStorage == null) {
+            unsegmentedStorage = new RocksDBKeyValueStorage(rocksDBConfiguration, metricsSystem);
+          }
+          return unsegmentedStorage;
+        }
+      case 1:
+        {
+          unsegmentedStorage = null;
+          if (segmentedStorage == null) {
+            segmentedStorage =
+                new RocksDBColumnarKeyValueStorage(rocksDBConfiguration, segments, metricsSystem);
+          }
+          return new SegmentedKeyValueStorageAdapter<>(segment, segmentedStorage);
+        }
+      default:
+        {
+          throw new IllegalStateException(
+              String.format(
+                  "Developer error: A supported database version (%d) was detected but there is no associated creation logic.",
+                  databaseVersion));
+        }
     }
   }
 
   protected Path storagePath(final BesuConfiguration commonConfiguration) {
     return commonConfiguration.getStoragePath();
   }
 
-  private boolean requiresInit() {
-    return segmentedStorage == null && unsegmentedStorage == null;
-  }
-
-  private void init(
-      final BesuConfiguration commonConfiguration, final MetricsSystem metricsSystem) {
+  private void init(final BesuConfiguration commonConfiguration) {
     try {
-      this.isSegmentIsolationSupported = databaseVersion(commonConfiguration) == DEFAULT_VERSION;
+      databaseVersion = readDatabaseVersion(commonConfiguration);
     } catch (final IOException e) {
       LOG.error("Failed to retrieve the RocksDB database meta version: {}", e.getMessage());
       throw new StorageException(e.getMessage(), e);
     }
-
-    final RocksDBConfiguration rocksDBConfiguration =
+    isSegmentIsolationSupported = databaseVersion >= 1;
+    rocksDBConfiguration =
         RocksDBConfigurationBuilder.from(configuration.get())
             .databaseDir(storagePath(commonConfiguration))
             .build();
+  }
 
-    if (isSegmentIsolationSupported) {
-      this.unsegmentedStorage = null;
-      this.segmentedStorage =
-          new RocksDBColumnarKeyValueStorage(rocksDBConfiguration, segments, metricsSystem);
-    } else {
-      this.unsegmentedStorage = new RocksDBKeyValueStorage(rocksDBConfiguration, metricsSystem);
-      this.segmentedStorage = null;
-    }
+  private boolean requiresInit() {
+    return segmentedStorage == null && unsegmentedStorage == null;
   }
 
-  private int databaseVersion(final BesuConfiguration commonConfiguration) throws IOException {
-    final Path databaseDir = storagePath(commonConfiguration);
+  private int readDatabaseVersion(final BesuConfiguration commonConfiguration) throws IOException {
+    final Path databaseDir = commonConfiguration.getStoragePath();
     final boolean databaseExists = databaseDir.resolve("IDENTITY").toFile().exists();
     final int databaseVersion;
     if (databaseExists) {
@@ -140,12 +158,29 @@ private int databaseVersion(final BesuConfiguration commonConfiguration) throws
       new DatabaseMetadata(databaseVersion).writeToDirectory(databaseDir);
     }
 
-    if (!SUPPORTED_VERSION.contains(databaseVersion)) {
+    if (!SUPPORTED_VERSIONS.contains(databaseVersion)) {
       final String message = "Unsupported RocksDB Metadata version of: " + databaseVersion;
       LOG.error(message);
       throw new StorageException(message);
     }
 
     return databaseVersion;
   }
+
+  @Override
+  public void close() throws IOException {
+    if (unsegmentedStorage != null) {
+      unsegmentedStorage.close();
+    }
+    if (segmentedStorage != null) {
+      segmentedStorage.close();
+    }
+  }
+
+  @Override
+  public boolean isSegmentIsolationSupported() {
+    return checkNotNull(
+        isSegmentIsolationSupported,
+        "Whether segment isolation is supported will be determined during creation. Call a creation method first");
+  }
 }

plugins/rocksdb/src/test/java/org/hyperledger/besu/plugin/services/storage/rocksdb/RocksDBKeyValueStorageFactoryTest.java

@@ -13,6 +13,7 @@
 package org.hyperledger.besu.plugin.services.storage.rocksdb;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.Mockito.when;
 
@@ -96,6 +97,22 @@ public void shouldDetectCorrectVersionIfMetadataFileExists() throws Exception {
     assertThat(storageFactory.isSegmentIsolationSupported()).isTrue();
   }
 
+  @Test
+  public void shouldDetectCorrectVersionInCaseOfRollback() throws Exception {
+    final Path tempDatabaseDir = temporaryFolder.newFolder().toPath().resolve("db");
+    Files.createDirectories(tempDatabaseDir);
+    when(commonConfiguration.getStoragePath()).thenReturn(tempDatabaseDir);
+    final RocksDBKeyValueStorageFactory storageFactory =
+        new RocksDBKeyValueStorageFactory(() -> rocksDbConfiguration, segments, 1);
+
+    storageFactory.create(segment, commonConfiguration, metricsSystem);
+    storageFactory.close();
+
+    final RocksDBKeyValueStorageFactory rolledbackStorageFactory =
+        new RocksDBKeyValueStorageFactory(() -> rocksDbConfiguration, segments, 0);
+    rolledbackStorageFactory.create(segment, commonConfiguration, metricsSystem);
+  }
+
   @Test
   public void shouldThrowExceptionWhenVersionNumberIsInvalid() throws Exception {
     final Path tempDatabaseDir = temporaryFolder.newFolder().toPath().resolve("db");
@@ -111,6 +128,17 @@ public void shouldThrowExceptionWhenVersionNumberIsInvalid() throws Exception {
         .isInstanceOf(StorageException.class);
   }
 
+  @Test
+  public void shouldSetSegmentationFieldDuringCreation() throws Exception {
+    final Path tempDatabaseDir = temporaryFolder.newFolder().toPath().resolve("db");
+    Files.createDirectories(tempDatabaseDir);
+    when(commonConfiguration.getStoragePath()).thenReturn(tempDatabaseDir);
+    final RocksDBKeyValueStorageFactory storageFactory =
+        new RocksDBKeyValueStorageFactory(() -> rocksDbConfiguration, segments);
+    storageFactory.create(segment, commonConfiguration, metricsSystem);
+    assertThatCode(storageFactory::isSegmentIsolationSupported).doesNotThrowAnyException();
+  }
+
   @Test
   public void shouldThrowExceptionWhenMetaDataFileIsCorrupted() throws Exception {
     final Path tempDatabaseDir = temporaryFolder.newFolder().toPath().resolve("db");