[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;
   }

clients/da-vinci-client/src/main/java/com/linkedin/davinci/storage/StorageService.java

@@ -77,6 +77,7 @@ public class StorageService extends AbstractVeniceService {
    * @param restoreDataPartitions indicates if store data needs to be restored.
    * @param restoreMetadataPartitions indicates if meta data needs to be restored.
    * @param checkWhetherStorageEngineShouldBeKeptOrNot check whether the local storage engine should be kept or not.
+   * @param checkWhetherStoragePartitionsShouldBeKeptOrNot check whether the partition is assigned and thus should be kept or not.
    */
   StorageService(
       VeniceConfigLoader configLoader,
@@ -88,6 +89,7 @@ public class StorageService extends AbstractVeniceService {
       boolean restoreDataPartitions,
       boolean restoreMetadataPartitions,
       Function<String, Boolean> checkWhetherStorageEngineShouldBeKeptOrNot,
+      Function<AbstractStorageEngine, Void> checkWhetherStoragePartitionsShouldBeKeptOrNot,
       Optional<Map<PersistenceType, StorageEngineFactory>> persistenceTypeToStorageEngineFactoryMapOptional) {
     String dataPath = configLoader.getVeniceServerConfig().getDataBasePath();
     if (!Utils.directoryExists(dataPath)) {
@@ -122,10 +124,97 @@ public class StorageService extends AbstractVeniceService {
           configLoader,
           restoreDataPartitions,
           restoreMetadataPartitions,
-          checkWhetherStorageEngineShouldBeKeptOrNot);
+          checkWhetherStorageEngineShouldBeKeptOrNot,
+          checkWhetherStoragePartitionsShouldBeKeptOrNot);
     }
   }
 
+  /**
+   * Allocates a new {@code StorageService} object.
+   * @param configLoader a config loader to load configs related to cluster and server.
+   * @param storageEngineStats storage engine related stats.
+   * @param rocksDBMemoryStats RocksDB memory consumption stats.
+   * @param storeVersionStateSerializer serializer for translating a store-version level state into avro-format.
+   * @param partitionStateSerializer serializer for translating a partition state into avro-format.
+   * @param storeRepository supports readonly operations to access stores
+   * @param restoreDataPartitions indicates if store data needs to be restored.
+   * @param restoreMetadataPartitions indicates if meta data needs to be restored.
+   * @param checkWhetherStorageEngineShouldBeKeptOrNot check whether the local storage engine should be kept or not.
+   */
+  StorageService(
+      VeniceConfigLoader configLoader,
+      AggVersionedStorageEngineStats storageEngineStats,
+      RocksDBMemoryStats rocksDBMemoryStats,
+      InternalAvroSpecificSerializer<StoreVersionState> storeVersionStateSerializer,
+      InternalAvroSpecificSerializer<PartitionState> partitionStateSerializer,
+      ReadOnlyStoreRepository storeRepository,
+      boolean restoreDataPartitions,
+      boolean restoreMetadataPartitions,
+      Function<String, Boolean> checkWhetherStorageEngineShouldBeKeptOrNot,
+      Optional<Map<PersistenceType, StorageEngineFactory>> persistenceTypeToStorageEngineFactoryMapOptional) {
+    String dataPath = configLoader.getVeniceServerConfig().getDataBasePath();
+    if (!Utils.directoryExists(dataPath)) {
+      if (!configLoader.getVeniceServerConfig().isAutoCreateDataPath()) {
+        throw new VeniceException(
+            "Data directory '" + dataPath + "' does not exist and " + ConfigKeys.AUTOCREATE_DATA_PATH
+                + " is disabled.");
+      }
+
+      File dataDir = new File(dataPath);
+      LOGGER.info("Creating data directory {}", dataDir.getAbsolutePath());
+      dataDir.mkdirs();
+    }
+
+    this.configLoader = configLoader;
+    this.serverConfig = configLoader.getVeniceServerConfig();
+    this.storageEngineRepository = new StorageEngineRepository();
+
+    this.aggVersionedStorageEngineStats = storageEngineStats;
+    this.rocksDBMemoryStats = rocksDBMemoryStats;
+    this.storeVersionStateSerializer = storeVersionStateSerializer;
+    this.partitionStateSerializer = partitionStateSerializer;
+    this.storeRepository = storeRepository;
+    if (persistenceTypeToStorageEngineFactoryMapOptional.isPresent()) {
+      this.persistenceTypeToStorageEngineFactoryMap = persistenceTypeToStorageEngineFactoryMapOptional.get();
+    } else {
+      this.persistenceTypeToStorageEngineFactoryMap = new HashMap<>();
+      initInternalStorageEngineFactories();
+    }
+    if (restoreDataPartitions || restoreMetadataPartitions) {
+      restoreAllStores(
+          configLoader,
+          restoreDataPartitions,
+          restoreMetadataPartitions,
+          checkWhetherStorageEngineShouldBeKeptOrNot,
+          se -> null);
+    }
+  }
+
+  public StorageService(
+      VeniceConfigLoader configLoader,
+      AggVersionedStorageEngineStats storageEngineStats,
+      RocksDBMemoryStats rocksDBMemoryStats,
+      InternalAvroSpecificSerializer<StoreVersionState> storeVersionStateSerializer,
+      InternalAvroSpecificSerializer<PartitionState> partitionStateSerializer,
+      ReadOnlyStoreRepository storeRepository,
+      boolean restoreDataPartitions,
+      boolean restoreMetadataPartitions,
+      Function<String, Boolean> checkWhetherStorageEngineShouldBeKeptOrNot,
+      Function<AbstractStorageEngine, Void> checkWhetherStoragePartitionsShouldBeKeptOrNot) {
+    this(
+        configLoader,
+        storageEngineStats,
+        rocksDBMemoryStats,
+        storeVersionStateSerializer,
+        partitionStateSerializer,
+        storeRepository,
+        restoreDataPartitions,
+        restoreMetadataPartitions,
+        checkWhetherStorageEngineShouldBeKeptOrNot,
+        checkWhetherStoragePartitionsShouldBeKeptOrNot,
+        Optional.empty());
+  }
+
   public StorageService(
       VeniceConfigLoader configLoader,
       AggVersionedStorageEngineStats storageEngineStats,
@@ -233,7 +322,8 @@ private void restoreAllStores(
       VeniceConfigLoader configLoader,
       boolean restoreDataPartitions,
       boolean restoreMetadataPartitions,
-      Function<String, Boolean> checkWhetherStorageEngineShouldBeKeptOrNot) {
+      Function<String, Boolean> checkWhetherStorageEngineShouldBeKeptOrNot,
+      Function<AbstractStorageEngine, Void> checkWhetherStoragePartitionsShouldBeKeptOrNot) {
     LOGGER.info("Start restoring all the stores persisted previously");
     for (Map.Entry<PersistenceType, StorageEngineFactory> entry: persistenceTypeToStorageEngineFactoryMap.entrySet()) {
       PersistenceType pType = entry.getKey();
@@ -254,6 +344,7 @@ private void restoreAllStores(
         if (checkWhetherStorageEngineShouldBeKeptOrNot.apply(storeName)) {
           try {
             storageEngine = openStore(storeConfig, () -> null);
+            checkWhetherStoragePartitionsShouldBeKeptOrNot.apply(storageEngine);
           } catch (Exception e) {
             if (ExceptionUtils.recursiveClassEquals(e, RocksDBException.class)) {
               LOGGER.warn("Encountered RocksDB error while opening store: {}", storeName, e);
@@ -467,7 +558,7 @@ public synchronized void closeStorageEngine(String kafkaTopic) {
   public void cleanupAllStores(VeniceConfigLoader configLoader) {
     // Load local storage and delete them safely.
     // TODO Just clean the data dir in case loading and deleting is too slow.
-    restoreAllStores(configLoader, true, true, s -> true);
+    restoreAllStores(configLoader, true, true, s -> true, se -> null);
     LOGGER.info("Start cleaning up all the stores persisted previously");
     storageEngineRepository.getAllLocalStorageEngines().stream().forEach(storageEngine -> {
       String storeName = storageEngine.getStoreVersionName();

clients/da-vinci-client/src/test/java/com/linkedin/davinci/storage/StorageServiceTest.java

@@ -115,6 +115,7 @@ public void testGetStoreAndUserPartitionsMapping() {
         true,
         true,
         (s) -> true,
+        (se) -> null,
         Optional.of(persistenceTypeToStorageEngineFactoryMap));
 
     Map<String, Set<Integer>> expectedMapping = new HashMap<>();

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,21 @@ 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(),
+        functionToCheckWhetherStoragePartitionsShouldBeKeptOrNot());
     storageEngineMetadataService =
         new StorageEngineMetadataService(storageService.getStorageEngineRepository(), partitionStateSerializer);
     services.add(storageEngineMetadataService);
@@ -693,6 +708,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;
   }