[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

@@ -21,6 +21,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.cleaner.BackupVersionOptimizationService;
@@ -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);
@@ -586,6 +600,10 @@ public void start() throws VeniceException {
       service.start();
     }
 
+    for (AbstractStorageEngine storageEngine: storageService.getStorageEngineRepository().getAllLocalStorageEngines()) {
+      functionToCheckWhetherStoragePartitionsShouldBeKeptOrNot().apply(storageEngine);
+    }
+
     for (ServiceDiscoveryAnnouncer serviceDiscoveryAnnouncer: serviceDiscoveryAnnouncers) {
       LOGGER.info("Registering to service discovery: {}", serviceDiscoveryAnnouncer);
       serviceDiscoveryAnnouncer.register();
@@ -693,6 +711,40 @@ protected VeniceConfigLoader getConfigLoader() {
     return veniceConfigLoader;
   }
 
+  protected final boolean isIsolatedIngestion() {
+    return veniceConfigLoader.getVeniceServerConfig().getIngestionMode().equals(IngestionMode.ISOLATED);
+  }
+
+  private Function<String, Boolean> functionToCheckWhetherStorageEngineShouldBeKeptOrNot() {
+    return storageEngineName -> true;
+  }
+
+  private Function<AbstractStorageEngine, Void> functionToCheckWhetherStoragePartitionsShouldBeKeptOrNot() {
+    return storageEngine -> {
+      String storageEngineName = storageEngine.toString();
+      String storeName = Version.parseStoreFromKafkaTopicName(storageEngineName);
+      PropertyKey.Builder propertyKeyBuilder =
+          new PropertyKey.Builder(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 null;
+    };
+  }
+
   public MetricsRepository getMetricsRepository() {
     return metricsRepository;
   }