[venice-server] Dropping unassigned partitions

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

@@ -77,7 +77,6 @@ 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,
@@ -89,7 +88,6 @@ 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)) {
@@ -124,97 +122,10 @@ public class StorageService extends AbstractVeniceService {
           configLoader,
           restoreDataPartitions,
           restoreMetadataPartitions,
-          checkWhetherStorageEngineShouldBeKeptOrNot,
-          checkWhetherStoragePartitionsShouldBeKeptOrNot);
+          checkWhetherStorageEngineShouldBeKeptOrNot);
     }
   }
 
-  /**
-   * 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,
@@ -322,8 +233,7 @@ private void restoreAllStores(
       VeniceConfigLoader configLoader,
       boolean restoreDataPartitions,
       boolean restoreMetadataPartitions,
-      Function<String, Boolean> checkWhetherStorageEngineShouldBeKeptOrNot,
-      Function<AbstractStorageEngine, Void> checkWhetherStoragePartitionsShouldBeKeptOrNot) {
+      Function<String, Boolean> checkWhetherStorageEngineShouldBeKeptOrNot) {
     LOGGER.info("Start restoring all the stores persisted previously");
     for (Map.Entry<PersistenceType, StorageEngineFactory> entry: persistenceTypeToStorageEngineFactoryMap.entrySet()) {
       PersistenceType pType = entry.getKey();
@@ -344,7 +254,6 @@ 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);
@@ -558,7 +467,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, se -> null);
+    restoreAllStores(configLoader, true, true, s -> true);
     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,7 +115,6 @@ 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

@@ -331,8 +331,7 @@ private List<AbstractVeniceService> createServices() {
         metadataRepo,
         whetherToRestoreDataPartitions,
         true,
-        functionToCheckWhetherStorageEngineShouldBeKeptOrNot(),
-        functionToCheckWhetherStoragePartitionsShouldBeKeptOrNot());
+        functionToCheckWhetherStorageEngineShouldBeKeptOrNot());
     storageEngineMetadataService =
         new StorageEngineMetadataService(storageService.getStorageEngineRepository(), partitionStateSerializer);
     services.add(storageEngineMetadataService);
@@ -601,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();