[server][dvc] Drop partitions asynchronously

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

@@ -252,7 +252,7 @@ public DaVinciBackend(
           .map(objectCacheConfig -> new ObjectCacheBackend(clientConfig, objectCacheConfig, schemaRepository));
 
       ingestionService = new KafkaStoreIngestionService(
-          storageService.getStorageEngineRepository(),
+          storageService,
           configLoader,
           storageMetadataService,
           clusterInfoProvider,

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

@@ -145,8 +145,7 @@ public void dropStoragePartitionGracefully(
     final int waitIntervalInSecond = 1;
     final int maxRetry = timeoutInSeconds / waitIntervalInSecond;
     getStoreIngestionService().stopConsumptionAndWait(storeConfig, partition, waitIntervalInSecond, maxRetry, true);
-    // Drops corresponding data partition from storage.
-    this.storageService.dropStorePartition(storeConfig, partition, removeEmptyStorageEngine);
+    getStoreIngestionService().dropStoragePartitionGracefully(storeConfig, partition);
   }
 
   @Override

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

@@ -695,7 +695,7 @@ private void initializeIsolatedIngestionServer() {
 
     // Create KafkaStoreIngestionService
     storeIngestionService = new KafkaStoreIngestionService(
-        storageService.getStorageEngineRepository(),
+        storageService,
         configLoader,
         storageMetadataService,
         clusterInfoProvider,

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/ActiveActiveStoreIngestionTask.java

@@ -15,6 +15,7 @@
 import com.linkedin.davinci.replication.merge.RmdSerDe;
 import com.linkedin.davinci.replication.merge.StringAnnotatedStoreSchemaCache;
 import com.linkedin.davinci.stats.AggVersionedIngestionStats;
+import com.linkedin.davinci.storage.StorageService;
 import com.linkedin.davinci.storage.chunking.ChunkedValueManifestContainer;
 import com.linkedin.davinci.storage.chunking.RawBytesChunkingAdapter;
 import com.linkedin.davinci.storage.chunking.SingleGetChunkingAdapter;
@@ -103,6 +104,7 @@ private static class ReusableObjects {
   private final ThreadLocal<ReusableObjects> threadLocalReusableObjects = ThreadLocal.withInitial(ReusableObjects::new);
 
   public ActiveActiveStoreIngestionTask(
+      StorageService storageService,
       StoreIngestionTaskFactory.Builder builder,
       Store store,
       Version version,
@@ -114,6 +116,7 @@ public ActiveActiveStoreIngestionTask(
       Optional<ObjectCacheBackend> cacheBackend,
       DaVinciRecordTransformerFunctionalInterface recordTransformerFunction) {
     super(
+        storageService,
         builder,
         store,
         version,

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/ConsumerActionType.java

@@ -4,7 +4,7 @@
  * An Enum enumerating all valid types of {@link ConsumerAction}.
  */
 public enum ConsumerActionType {
-  SUBSCRIBE(1), UNSUBSCRIBE(1), RESET_OFFSET(1), PAUSE(1), RESUME(1),
+  SUBSCRIBE(1), UNSUBSCRIBE(1), RESET_OFFSET(1), PAUSE(1), RESUME(1), DROP_PARTITION(1),
   /**
    * KILL action has higher priority than others, so that once KILL action is added to the action queue,
    * we will process it immediately to avoid doing throw-away works.

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/KafkaStoreIngestionService.java

@@ -29,8 +29,8 @@
 import com.linkedin.davinci.stats.AggVersionedIngestionStats;
 import com.linkedin.davinci.stats.ParticipantStoreConsumptionStats;
 import com.linkedin.davinci.stats.ingestion.heartbeat.HeartbeatMonitoringService;
-import com.linkedin.davinci.storage.StorageEngineRepository;
 import com.linkedin.davinci.storage.StorageMetadataService;
+import com.linkedin.davinci.storage.StorageService;
 import com.linkedin.davinci.store.cache.backend.ObjectCacheBackend;
 import com.linkedin.davinci.store.view.VeniceViewWriterFactory;
 import com.linkedin.venice.SSLConfig;
@@ -130,6 +130,7 @@ public class KafkaStoreIngestionService extends AbstractVeniceService implements
   private static final String GROUP_ID_FORMAT = "%s_%s";
 
   private static final Logger LOGGER = LogManager.getLogger(KafkaStoreIngestionService.class);
+  private final StorageService storageService;
 
   private final VeniceConfigLoader veniceConfigLoader;
 
@@ -190,7 +191,7 @@ public class KafkaStoreIngestionService extends AbstractVeniceService implements
   private final ExecutorService aaWCWorkLoadProcessingThreadPool;
 
   public KafkaStoreIngestionService(
-      StorageEngineRepository storageEngineRepository,
+      StorageService storageService,
       VeniceConfigLoader veniceConfigLoader,
       StorageMetadataService storageMetadataService,
       ClusterInfoProvider clusterInfoProvider,
@@ -212,6 +213,7 @@ public KafkaStoreIngestionService(
       PubSubClientsFactory pubSubClientsFactory,
       Optional<SSLFactory> sslFactory,
       HeartbeatMonitoringService heartbeatMonitoringService) {
+    this.storageService = storageService;
     this.cacheBackend = cacheBackend;
     this.recordTransformerFunction = recordTransformerFunction;
     this.storageMetadataService = storageMetadataService;
@@ -448,7 +450,7 @@ public void handleStoreDeleted(Store store) {
 
     ingestionTaskFactory = StoreIngestionTaskFactory.builder()
         .setVeniceWriterFactory(veniceWriterFactory)
-        .setStorageEngineRepository(storageEngineRepository)
+        .setStorageEngineRepository(storageService.getStorageEngineRepository())
         .setStorageMetadataService(storageMetadataService)
         .setLeaderFollowerNotifiersQueue(leaderFollowerNotifiers)
         .setSchemaRepository(schemaRepo)
@@ -519,6 +521,7 @@ private StoreIngestionTask createStoreIngestionTask(
     };
 
     return ingestionTaskFactory.getNewIngestionTask(
+        storageService,
         store,
         version,
         getKafkaConsumerProperties(veniceStoreVersionConfig),
@@ -920,6 +923,37 @@ public void stopConsumptionAndWait(
     }
   }
 
+  /**
+   * Drops the corresponding Venice Partition gracefully.
+   * This should only be called after {@link #stopConsumptionAndWait} has been called
+   * @param veniceStore Venice Store for the partition.
+   * @param partitionId Venice partition's id.
+   */
+  public void dropStoragePartitionGracefully(VeniceStoreVersionConfig veniceStore, int partitionId) {
+    final String topic = veniceStore.getStoreVersionName();
+
+    if (isPartitionConsuming(topic, partitionId)) {
+      throw new VeniceException("Tried to drop storage partition that is still consuming");
+    }
+
+    try (AutoCloseableLock ignore = topicLockManager.getLockForResource(topic)) {
+      StoreIngestionTask ingestionTask = topicNameToIngestionTaskMap.get(topic);
+      if (ingestionTask != null && ingestionTask.isRunning()) {
+        LOGGER.info(
+            "Ingestion task is still running for Topic {}. Dropping partition {} asynchronously",
+            veniceStore.getStoreVersionName(),
+            partitionId);
+        ingestionTask.dropPartition(new PubSubTopicPartitionImpl(pubSubTopicRepository.getTopic(topic), partitionId));
+      } else {
+        LOGGER.info(
+            "Ingestion task isn't running for Topic {}. Dropping partition {} synchronously",
+            veniceStore.getStoreVersionName(),
+            partitionId);
+        this.storageService.dropStorePartition(veniceStore, partitionId, true);
+      }
+    }
+  }
+
   /**
    * This function will try to kill the ingestion tasks belonging to non-current versions.
    * And this is mainly being used by memory limiter feature to free up resources when encountering memory

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/LeaderFollowerStoreIngestionTask.java

@@ -23,6 +23,7 @@
 import com.linkedin.davinci.schema.merge.CollectionTimestampMergeRecordHelper;
 import com.linkedin.davinci.schema.merge.MergeRecordHelper;
 import com.linkedin.davinci.stats.ingestion.heartbeat.HeartbeatMonitoringService;
+import com.linkedin.davinci.storage.StorageService;
 import com.linkedin.davinci.storage.chunking.ChunkedValueManifestContainer;
 import com.linkedin.davinci.storage.chunking.GenericRecordChunkingAdapter;
 import com.linkedin.davinci.store.AbstractStorageEngine;
@@ -204,6 +205,7 @@ public class LeaderFollowerStoreIngestionTask extends StoreIngestionTask {
   private final Version version;
 
   public LeaderFollowerStoreIngestionTask(
+      StorageService storageService,
       StoreIngestionTaskFactory.Builder builder,
       Store store,
       Version version,
@@ -215,6 +217,7 @@ public LeaderFollowerStoreIngestionTask(
       Optional<ObjectCacheBackend> cacheBackend,
       DaVinciRecordTransformerFunctionalInterface recordTransformerFunction) {
     super(
+        storageService,
         builder,
         store,
         version,

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTask.java

@@ -2,6 +2,7 @@
 
 import static com.linkedin.davinci.ingestion.LagType.OFFSET_LAG;
 import static com.linkedin.davinci.ingestion.LagType.TIME_LAG;
+import static com.linkedin.davinci.kafka.consumer.ConsumerActionType.DROP_PARTITION;
 import static com.linkedin.davinci.kafka.consumer.ConsumerActionType.RESET_OFFSET;
 import static com.linkedin.davinci.kafka.consumer.ConsumerActionType.SUBSCRIBE;
 import static com.linkedin.davinci.kafka.consumer.ConsumerActionType.UNSUBSCRIBE;
@@ -34,6 +35,7 @@
 import com.linkedin.davinci.stats.HostLevelIngestionStats;
 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.davinci.store.StoragePartitionConfig;
 import com.linkedin.davinci.store.cache.backend.ObjectCacheBackend;
@@ -189,6 +191,7 @@ public abstract class StoreIngestionTask implements Runnable, Closeable {
       Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2);
 
   /** storage destination for consumption */
+  protected final StorageService storageService;
   protected final StorageEngineRepository storageEngineRepository;
   protected final AbstractStorageEngine storageEngine;
 
@@ -342,8 +345,10 @@ public abstract class StoreIngestionTask implements Runnable, Closeable {
   protected final boolean batchReportIncPushStatusEnabled;
 
   protected final ExecutorService parallelProcessingThreadPool;
+  private final Set<Integer> pendingPartitionDrops = Collections.synchronizedSet(new HashSet<>());
 
   public StoreIngestionTask(
+      StorageService storageService,
       StoreIngestionTaskFactory.Builder builder,
       Store store,
       Version version,
@@ -361,6 +366,7 @@ public StoreIngestionTask(
     this.databaseSyncBytesIntervalForTransactionalMode = storeConfig.getDatabaseSyncBytesIntervalForTransactionalMode();
     this.databaseSyncBytesIntervalForDeferredWriteMode = storeConfig.getDatabaseSyncBytesIntervalForDeferredWriteMode();
     this.kafkaProps = kafkaConsumerProperties;
+    this.storageService = storageService;
     this.storageEngineRepository = builder.getStorageEngineRepository();
     this.storageMetadataService = builder.getStorageMetadataService();
     this.storeRepository = builder.getMetadataRepo();
@@ -630,6 +636,18 @@ public synchronized CompletableFuture<Void> unSubscribePartition(
     return consumerAction.getFuture();
   }
 
+  /**
+   *
+   * Adds an asynchronous partition drop request for the task.
+   * This is always a Helix triggered action
+   */
+  public synchronized void dropPartition(PubSubTopicPartition topicPartition) {
+    throwIfNotRunning();
+    ConsumerAction consumerAction = new ConsumerAction(DROP_PARTITION, topicPartition, nextSeqNum(), true);
+    consumerActionsQueue.add(consumerAction);
+    pendingPartitionDrops.add(topicPartition.getPartitionNumber());
+  }
+
   public boolean hasAnySubscription() {
     return !partitionConsumptionStateMap.isEmpty();
   }
@@ -2133,8 +2151,27 @@ protected void processCommonConsumerAction(ConsumerAction consumerAction) throws
         break;
       case KILL:
         LOGGER.info("Kill this consumer task for Topic: {}", topic);
+
+        if (!this.pendingPartitionDrops.isEmpty()) {
+          LOGGER.info(
+              "Partitions {} are pending to be dropped for Topic: {}. Dropping them before killing consumer task.",
+              this.pendingPartitionDrops,
+              topic);
+          synchronized (this.pendingPartitionDrops) {
+            for (Integer partitionToDrop: this.pendingPartitionDrops) {
+              this.storageService.dropStorePartition(storeConfig, partitionToDrop, true);
+              this.pendingPartitionDrops.remove(partitionToDrop);
+            }
+          }
+        }
         // Throw the exception here to break the consumption loop, and then this task is marked as error status.
         throw new VeniceIngestionTaskKilledException(KILLED_JOB_MESSAGE + topic);
+      case DROP_PARTITION:
+        LOGGER.info("{} Dropping partition: {}", ingestionTaskName, topicPartition);
+        this.storageService.dropStorePartition(storeConfig, partition, true);
+        this.pendingPartitionDrops.remove(partition);
+        LOGGER.info("{} Dropped partition: {}", ingestionTaskName, topicPartition);
+        break;
       default:
         throw new UnsupportedOperationException(operation.name() + " is not supported in " + getClass().getName());
     }

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTaskFactory.java

@@ -11,6 +11,7 @@
 import com.linkedin.davinci.stats.ingestion.heartbeat.HeartbeatMonitoringService;
 import com.linkedin.davinci.storage.StorageEngineRepository;
 import com.linkedin.davinci.storage.StorageMetadataService;
+import com.linkedin.davinci.storage.StorageService;
 import com.linkedin.davinci.store.cache.backend.ObjectCacheBackend;
 import com.linkedin.davinci.store.view.VeniceViewWriterFactory;
 import com.linkedin.venice.kafka.protocol.state.PartitionState;
@@ -43,6 +44,7 @@ private StoreIngestionTaskFactory(Builder builder) {
   }
 
   public StoreIngestionTask getNewIngestionTask(
+      StorageService storageService,
       Store store,
       Version version,
       Properties kafkaConsumerProperties,
@@ -54,6 +56,7 @@ public StoreIngestionTask getNewIngestionTask(
       DaVinciRecordTransformerFunctionalInterface recordTransformerFunction) {
     if (version.isActiveActiveReplicationEnabled()) {
       return new ActiveActiveStoreIngestionTask(
+          storageService,
           builder,
           store,
           version,
@@ -66,6 +69,7 @@ public StoreIngestionTask getNewIngestionTask(
           recordTransformerFunction);
     }
     return new LeaderFollowerStoreIngestionTask(
+        storageService,
         builder,
         store,
         version,

clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/ActiveActiveStoreIngestionTaskTest.java

@@ -28,6 +28,7 @@
 import com.linkedin.davinci.stats.AggVersionedIngestionStats;
 import com.linkedin.davinci.stats.HostLevelIngestionStats;
 import com.linkedin.davinci.storage.StorageEngineRepository;
+import com.linkedin.davinci.storage.StorageService;
 import com.linkedin.davinci.storage.chunking.ChunkedValueManifestContainer;
 import com.linkedin.davinci.storage.chunking.ChunkingUtils;
 import com.linkedin.davinci.store.AbstractStorageEngine;
@@ -231,6 +232,7 @@ public void testisReadyToServeAnnouncedWithRTLag() {
     Version mockVersion = new VersionImpl(STORE_NAME, 1, PUSH_JOB_ID);
     mockVersion.setHybridStoreConfig(hybridStoreConfig);
 
+    StorageService storageService = mock(StorageService.class);
     Store store = new ZKStore(
         STORE_NAME,
         "Felix",
@@ -250,6 +252,7 @@ public void testisReadyToServeAnnouncedWithRTLag() {
     VeniceStoreVersionConfig storeVersionConfig =
         new VeniceStoreVersionConfig(STORE_NAME + "_v1", new VeniceProperties(kafkaConsumerProperties));
     ActiveActiveStoreIngestionTask ingestionTask = new ActiveActiveStoreIngestionTask(
+        storageService,
         builder,
         store,
         mockVersion,