[venice-server] Dropping unassigned partitions

clients/da-vinci-client/src/main/java/com/linkedin/davinci/helix/HelixParticipationService.java

@@ -427,6 +427,10 @@ public Instance getInstance() {
     return instance;
   }
 
+  public SafeHelixManager getHelixManager() {
+    return helixManager;
+  }
+
   public VeniceOfflinePushMonitorAccessor getVeniceOfflinePushMonitorAccessor() {
     return veniceOfflinePushMonitorAccessor;
   }

services/venice-server/src/main/java/com/linkedin/venice/server/VeniceServer.java

@@ -19,6 +19,7 @@
 import com.linkedin.davinci.storage.StorageEngineRepository;
 import com.linkedin.davinci.storage.StorageMetadataService;
 import com.linkedin.davinci.storage.StorageService;
+import com.linkedin.davinci.store.AbstractStorageEngine;
 import com.linkedin.venice.acl.DynamicAccessController;
 import com.linkedin.venice.acl.StaticAccessController;
 import com.linkedin.venice.blobtransfer.BlobTransferManager;
@@ -42,10 +43,12 @@
 import com.linkedin.venice.listener.ServerReadMetadataRepository;
 import com.linkedin.venice.listener.ServerStoreAclHandler;
 import com.linkedin.venice.listener.StoreValueSchemasCacheService;
+import com.linkedin.venice.meta.IngestionMode;
 import com.linkedin.venice.meta.ReadOnlyLiveClusterConfigRepository;
 import com.linkedin.venice.meta.ReadOnlySchemaRepository;
 import com.linkedin.venice.meta.ReadOnlyStoreRepository;
 import com.linkedin.venice.meta.StaticClusterInfoProvider;
+import com.linkedin.venice.meta.Version;
 import com.linkedin.venice.pubsub.PubSubClientsFactory;
 import com.linkedin.venice.schema.SchemaReader;
 import com.linkedin.venice.security.SSLFactory;
@@ -66,11 +69,16 @@
 import io.tehuti.metrics.MetricsRepository;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Function;
+import org.apache.helix.PropertyKey;
+import org.apache.helix.model.IdealState;
 import org.apache.helix.zookeeper.impl.client.ZkClient;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -310,14 +318,20 @@ private List<AbstractVeniceService> createServices() {
         ? new RocksDBMemoryStats(metricsRepository, "RocksDBMemoryStats", plainTableEnabled)
         : null;
 
+    boolean whetherToRestoreDataPartitions = !isIsolatedIngestion()
+        || veniceConfigLoader.getVeniceServerConfig().freezeIngestionIfReadyToServeOrLocalDataExists();
+
     // Create and add StorageService. storeRepository will be populated by StorageService
     storageService = new StorageService(
         veniceConfigLoader,
         storageEngineStats,
         rocksDBMemoryStats,
         storeVersionStateSerializer,
         partitionStateSerializer,
-        metadataRepo);
+        metadataRepo,
+        whetherToRestoreDataPartitions,
+        true,
+        functionToCheckWhetherStorageEngineShouldBeKeptOrNot());
     storageEngineMetadataService =
         new StorageEngineMetadataService(storageService.getStorageEngineRepository(), partitionStateSerializer);
     services.add(storageEngineMetadataService);
@@ -692,6 +706,45 @@ protected VeniceConfigLoader getConfigLoader() {
     return veniceConfigLoader;
   }
 
+  protected final boolean isIsolatedIngestion() {
+    return veniceConfigLoader.getVeniceServerConfig().getIngestionMode().equals(IngestionMode.ISOLATED);
+  }
+
+  private Function<String, Boolean> functionToCheckWhetherStorageEngineShouldBeKeptOrNot() {
+    return storageEngineName -> {
+      String storeName = Version.parseStoreFromKafkaTopicName(storageEngineName);
+
+      StorageEngineRepository storageEngineRepository = new StorageEngineRepository();
+
+      AbstractStorageEngine storageEngine = storageEngineRepository.getLocalStorageEngine(storageEngineName);
+
+      if (storageEngine == null) {
+        return true;
+      }
+
+      PropertyKey.Builder propertyKeyBuilder =
+          new PropertyKey.Builder(this.veniceConfigLoader.getVeniceClusterConfig().getClusterName());
+      IdealState idealState = getHelixParticipationService().getHelixManager()
+          .getHelixDataAccessor()
+          .getProperty(propertyKeyBuilder.idealStates(storeName));
+
+      Set<Integer> idealStatePartitionIds = new HashSet<>();
+      idealState.getPartitionSet().stream().forEach(partitionId -> {
+        idealStatePartitionIds.add(Integer.parseInt(partitionId));
+      });
+      Set<Integer> storageEnginePartitionIds = storageEngine.getPartitionIds();
+
+      for (Integer storageEnginePartitionId: storageEnginePartitionIds) {
+        if (idealStatePartitionIds.contains(storageEnginePartitionId)) {
+          continue;
+        }
+        storageEngine.dropPartition(storageEnginePartitionId);
+      }
+
+      return true;
+    };
+  }
+
   public MetricsRepository getMetricsRepository() {
     return metricsRepository;
   }