[server] Restore data partition in isolated process in Venice Server (#810)

This PR tries to address the issue of cleaning up data correctly for ingestion isolation enabled server.
Originally, we observed that certain data partition is not deleted when a OFFLINE -> DROPPED message is issued for stale version upon server restart. The message is not executed as we did not restore all the data partitions upon server start in any of the processes.
The first fix we implemented is in #617 , where we added a clean up service in isolated process. However, it is not correct, as the clean up service tries to remove partitions opened in the forked process, which is metadata partition only. This leaves the data partitions still unopened and lingering forever.

This PR tries to fix it thoroughly by adjusting the data restore logic for II. It will by default restore everything in II process. For those partitions to receive OFFLINE->DROPPED message, it will be sent to isolated process and get executed. For old storage engines, it shall also be opened and should be cleaned up eventually, as restore metadata partition will recreate the metadata partition, even if it is deleted before.

Author: sixpluszero

clients/da-vinci-client/src/main/java/com/linkedin/davinci/ingestion/IsolatedIngestionBackend.java

@@ -390,15 +390,17 @@ void executeCommandWithRetry(
       Runnable mainProcessCommandRunnable) {
     boolean isTopicPartitionHosted = isTopicPartitionHosted(topicName, partition);
 
-    // drop storage partition in main process if it does not exist
-    if (command == REMOVE_PARTITION && !isTopicPartitionHosted) {
-      mainProcessCommandRunnable.run();
-      return;
-    }
-
     do {
+      /**
+       * There are two cases where we execute the command to the main process:
+       * (1) The main process metadata tells that the topic partition is hosted in main process.
+       * (2) The topic partition has never been associated with the server, and the incoming command is not START_CONSUMPTION
+       * or REMOVE_PARTITION. For these two commands, it should by default be sent to isolated process. Here it needs to
+       * make sure topic partition subscription request has never been issued to the server, so the check "isTopicPartitionHosted"
+       * should be false in this case.
+       */
       if (isTopicPartitionHostedInMainProcess(topicName, partition)
-          || !isTopicPartitionHosted && command != START_CONSUMPTION) {
+          || (!isTopicPartitionHosted && command != START_CONSUMPTION && command != REMOVE_PARTITION)) {
         LOGGER.info("Executing command {} of topic: {}, partition: {} in main process.", command, topicName, partition);
         mainProcessCommandRunnable.run();
         return;

clients/da-vinci-client/src/main/java/com/linkedin/davinci/ingestion/isolated/IsolatedIngestionServer.java

@@ -472,7 +472,7 @@ public boolean isResourceSubscribed(String topicName, int partition) {
     return subscription.get();
   }
 
-  public void maybeSubscribeNewResource(String topicName, int partition) {
+  public void maybeInitializeResourceHostingMetadata(String topicName, int partition) {
     getTopicPartitionSubscriptionMap().computeIfAbsent(topicName, s -> new VeniceConcurrentHashMap<>())
         .computeIfAbsent(partition, p -> new AtomicBoolean(true));
   }
@@ -715,27 +715,30 @@ private void initializeIsolatedIngestionServer() {
         metricsRepository,
         storeRepository,
         serverConfig.isUnregisterMetricForDeletedStoreEnabled());
+    boolean isDaVinciClient = veniceMetadataRepositoryBuilder.isDaVinciClient();
+
     /**
-     * The reason of not to restore the data partitions during initialization of storage service is:
-     * 1. During first fresh start up with no data on disk, we don't need to restore anything
-     * 2. During fresh start up with data on disk (aka bootstrap), we will receive messages to subscribe to the partition
+     * For isolated ingestion server, we always restore metadata partition, but we will only restore data partition in server,
+     * not Da Vinci.
+     * The reason of not to restore the data partitions during initialization of storage service in DVC is:
+     * 1. During fresh start up with data on disk (aka bootstrap), we will receive messages to subscribe to the partition
      * and it will re-open the partition on demand.
-     * 3. During crash recovery restart, partitions that are already ingestion will be opened by parent process and we
+     * 2. During crash recovery restart, partitions that are already ingestion will be opened by parent process and we
      * should not try to open it. The remaining ingestion tasks will open the storage engines.
-     * Also, we don't need to restore metadata partition here, as all subscribe requests sent to forked process will
-     * automatically open the metadata partitions. Also, during Da Vinci bootstrap, main process will need to open the
-     * metadata partition of storage engines in order to perform full cleanup of stale versions.
+     *
+     * The reason to restore data partitions during initialization of storage service in server is:
+     * 1. Helix can issue OFFLINE->DROPPED state transition for stale partition in anytime, including server restart time.
+     * If data partition is not restored here (as well as the main process), it will never be dropped so the RocksDB storage
+     * will be leaking.
      */
-    boolean isDaVinciClient = veniceMetadataRepositoryBuilder.isDaVinciClient();
-
     storageService = new StorageService(
         configLoader,
         storageEngineStats,
         rocksDBMemoryStats,
         storeVersionStateSerializer,
         partitionStateSerializer,
         storeRepository,
-        false,
+        !isDaVinciClient,
         true);
     storageService.start();
 

clients/da-vinci-client/src/main/java/com/linkedin/davinci/ingestion/isolated/IsolatedIngestionServerHandler.java

@@ -162,7 +162,7 @@ private IngestionTaskReport handleIngestionTaskCommand(IngestionTaskCommand inge
       storeConfig.setRestoreDataPartitions(false);
       switch (ingestionCommandType) {
         case START_CONSUMPTION:
-          isolatedIngestionServer.maybeSubscribeNewResource(topicName, partitionId);
+          isolatedIngestionServer.maybeInitializeResourceHostingMetadata(topicName, partitionId);
           validateAndExecuteCommand(ingestionCommandType, report, () -> {
             ReadOnlyStoreRepository storeRepository = isolatedIngestionServer.getStoreRepository();
             // For subscription based store repository, we will need to subscribe to the store explicitly.
@@ -201,6 +201,12 @@ private IngestionTaskReport handleIngestionTaskCommand(IngestionTaskCommand inge
           isolatedIngestionServer.cleanupTopicState(topicName);
           break;
         case REMOVE_PARTITION:
+          /**
+           * For this command, we will try to initialize the isolated process metadata for the topic partition, if the
+           * topic partition has never been associated with the server. This is needed as data partition is by default
+           * restored in the isolated process, and thus it should be able to be removed by Helix command.
+           */
+          isolatedIngestionServer.maybeInitializeResourceHostingMetadata(topicName, partitionId);
           validateAndExecuteCommand(ingestionCommandType, report, () -> {
             /**
              * Here we do not allow storage service to clean up "empty" storage engine. When ingestion isolation is turned on,