[server][da-vinci] Honor IngestionPauseMode by unsubscribing Kafka consumers (linkedin#2763)

Author: misyel

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

@@ -46,6 +46,7 @@
 import com.linkedin.venice.compression.CompressionStrategy;
 import com.linkedin.venice.exceptions.VeniceException;
 import com.linkedin.venice.exceptions.VeniceMessageException;
+import com.linkedin.venice.exceptions.VeniceNoStoreException;
 import com.linkedin.venice.exceptions.VeniceTimeoutException;
 import com.linkedin.venice.exceptions.validation.DuplicateDataException;
 import com.linkedin.venice.exceptions.validation.FatalDataValidationException;
@@ -620,6 +621,7 @@ protected void checkLongRunningTaskState() throws InterruptedException {
     boolean pushTimeout = false;
     Set<Integer> timeoutPartitions = null;
     long checkStartTimeInNS = System.nanoTime();
+    maybeTransitionPauseState();
     for (PartitionConsumptionState partitionConsumptionState: getPartitionConsumptionStateMap().values()) {
       final int partition = partitionConsumptionState.getPartition();
 
@@ -675,7 +677,7 @@ protected void checkLongRunningTaskState() throws InterruptedException {
        * online replica continue serving and do not close ingestion task.
        */
       if (!partitionConsumptionState.isComplete() && !partitionConsumptionState.isErrorReported()
-          && LatencyUtils.getElapsedTimeFromMsToMs(
+          && !partitionConsumptionState.isStoreLevelPaused() && LatencyUtils.getElapsedTimeFromMsToMs(
               partitionConsumptionState.getConsumptionStartTimeInMs()) > getBootstrapTimeoutInMs()) {
         if (!pushTimeout) {
           pushTimeout = true;
@@ -871,6 +873,143 @@ protected void checkLongRunningTaskState() throws InterruptedException {
         getVersionTopic().getName());
   }
 
+  /**
+   * Reconciles each PCS's store-level pause state against the current store metadata.
+   * <p>
+   * On a transition into paused: removes the partition's current leader topic (or VT for
+   * followers) from every per-broker consumer service via {@link #consumerUnSubscribeAllTopics}.
+   * Note: cross-region RT subscriptions for an A/A leader are managed by the leader's normal A/A
+   * code path, not by this hook — only the topic in the OffsetRecord is dropped here.
+   * On a transition out of paused: resubscribes via {@link #resubscribe}, which rewinds from the
+   * persisted offset to the leader topic recorded in the OffsetRecord (followers: local VT;
+   * leaders: the recorded leader topic).
+   * <p>
+   * The PCS pause flag is set to its target value <em>before</em> the long-running unsubscribe /
+   * resubscribe so the disk-quota no-op guard covers the entire transition window. Each PCS is
+   * processed inside a try/catch so a failure on one partition does not abandon the others.
+   */
+  void maybeTransitionPauseState() throws InterruptedException {
+    Store store;
+    try {
+      store = storeRepository.getStoreOrThrow(storeName);
+    } catch (VeniceNoStoreException e) {
+      // Store metadata is genuinely unavailable (deleted, not yet propagated, etc.). Leave pause
+      // state as-is and retry next iteration. Throttle WARN via the inherited filter so a
+      // continuously-failing lookup logs at most once per window instead of once per SIT
+      // iteration. Other RuntimeExceptions propagate so real bugs surface.
+      if (!REDUNDANT_LOGGING_FILTER.isRedundantException(storeName + "-maybeTransitionPauseState-store-not-found")) {
+        LOGGER.warn("Store {} not found while reconciling pause state; skipping transition.", storeName);
+      }
+      return;
+    }
+    boolean shouldPause = shouldPauseForStore(store);
+    boolean transitioned = false;
+    // Lazily probe the consumer subscription state once per loop (SIT-wide), only when a
+    // RECONCILE_FORCE_UNSUBSCRIBE could fire (shouldPause is set). consumerHasAnySubscription is
+    // a single fast call against aggKafkaConsumerService; consumerUnSubscribeAllTopics itself is
+    // self-gating per topic, so triggering reconcile when only some partitions are subscribed is
+    // benign (no-ops for the others).
+    boolean anySubscriptionForSit = shouldPause && consumerHasAnySubscription();
+    for (PartitionConsumptionState pcs: getPartitionConsumptionStateMap().values()) {
+      PauseStateTransition transition = decidePauseTransition(pcs, shouldPause, anySubscriptionForSit);
+      if (transition == PauseStateTransition.NO_CHANGE) {
+        continue;
+      }
+      try {
+        if (transition == PauseStateTransition.ENTER_PAUSE) {
+          // Flip the flag BEFORE the long-running unsubscribe so concurrent disk-quota callbacks
+          // see the new state and no-op for the entire window.
+          pcs.setStoreLevelPaused(true);
+          consumerUnSubscribeAllTopics(pcs);
+          LOGGER.info(
+              "Store-level pause activated for replica: {} — unsubscribed from Kafka",
+              Utils.getReplicaId(getKafkaVersionTopic(), pcs.getPartition()));
+        } else if (transition == PauseStateTransition.RECONCILE_FORCE_UNSUBSCRIBE) {
+          consumerUnSubscribeAllTopics(pcs);
+          LOGGER.info(
+              "Store-level pause re-applied for replica: {} — subscription was reattached, unsubscribed again",
+              Utils.getReplicaId(getKafkaVersionTopic(), pcs.getPartition()));
+        } else { // EXIT_PAUSE
+          // Resubscribe BEFORE clearing the flag so quota callbacks stay no-op until the consumer
+          // is back online; if resubscribe throws we leave the flag set so the next iteration
+          // retries instead of leaving the partition dark.
+          resubscribe(pcs);
+          pcs.setStoreLevelPaused(false);
+          pcs.resetConsumptionStartTimeInMs();
+          LOGGER.info(
+              "Store-level pause deactivated for replica: {} — resubscribed from persisted offset",
+              Utils.getReplicaId(getKafkaVersionTopic(), pcs.getPartition()));
+        }
+        transitioned = true;
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw e;
+      } catch (Exception e) {
+        LOGGER.error(
+            "Failed to apply pause transition {} for replica: {}; will retry next iteration.",
+            transition,
+            Utils.getReplicaId(getKafkaVersionTopic(), pcs.getPartition()),
+            e);
+      }
+    }
+    if (transitioned) {
+      // Reflect actual post-loop state, not intent — partial-failure cases shouldn't flip the
+      // gauge to 0 while some PCSes remain paused.
+      boolean anyPcsPaused = false;
+      for (PartitionConsumptionState pcs: getPartitionConsumptionStateMap().values()) {
+        if (pcs != null && pcs.isStoreLevelPaused()) {
+          anyPcsPaused = true;
+          break;
+        }
+      }
+      versionedIngestionStats.setStoreLevelPausedGauge(storeName, versionNumber, anyPcsPaused);
+    }
+  }
+
+  /**
+   * Result of evaluating a single PCS against the desired pause state.
+   * {@code RECONCILE_FORCE_UNSUBSCRIBE} fires when the PCS is already flagged paused but a
+   * subscription has crept back (e.g., a leader/follower transition or topic switch reattached
+   * the consumer); the reconcile loop force-unsubscribes again so ingestion can't resume while
+   * shouldPause is still true.
+   */
+  enum PauseStateTransition {
+    ENTER_PAUSE, EXIT_PAUSE, RECONCILE_FORCE_UNSUBSCRIBE, NO_CHANGE
+  }
+
+  /**
+   * Pure decision function: given the current PCS pause flag, the desired pause state, and
+   * whether the partition still has any active Kafka subscription, return the transition that
+   * needs to happen (if any). Side-effect-free and trivially unit-testable.
+   *
+   * <p>{@code RECONCILE_FORCE_UNSUBSCRIBE} fires when {@code shouldPause} is true and the PCS is
+   * already flagged paused but {@code hasAnyActiveSubscription} reports a live subscription —
+   * makes the reconcile loop idempotent against out-of-band re-subscriptions (L/F transitions,
+   * topic switches) that would otherwise leave a paused store ingesting.
+   *
+   * <p>Returns NO_CHANGE for a {@code null} PCS — defensive guard mirroring
+   * {@link StoreIngestionTask#shouldSkipQuotaCallbackForStoreLevelPause(PartitionConsumptionState)}.
+   */
+  static PauseStateTransition decidePauseTransition(
+      PartitionConsumptionState pcs,
+      boolean shouldPause,
+      boolean hasAnyActiveSubscription) {
+    if (pcs == null) {
+      return PauseStateTransition.NO_CHANGE;
+    }
+    boolean isPaused = pcs.isStoreLevelPaused();
+    if (shouldPause && !isPaused) {
+      return PauseStateTransition.ENTER_PAUSE;
+    }
+    if (shouldPause && isPaused && hasAnyActiveSubscription) {
+      return PauseStateTransition.RECONCILE_FORCE_UNSUBSCRIBE;
+    }
+    if (!shouldPause && pcs.isStoreLevelPaused()) {
+      return PauseStateTransition.EXIT_PAUSE;
+    }
+    return PauseStateTransition.NO_CHANGE;
+  }
+
   protected static boolean checkWhetherToCloseUnusedVeniceWriter(
       Lazy<VeniceWriter<byte[], byte[], byte[]>> veniceWriterLazy,
       Lazy<VeniceWriter<byte[], byte[], byte[]>> veniceWriterForRealTimeLazy,
@@ -3457,15 +3596,29 @@ protected void consumerBatchUnsubscribeAllTopics() {
    */
   @Override
   public void consumerUnSubscribeAllTopics(PartitionConsumptionState partitionConsumptionState) {
-    PubSubTopic leaderTopic = partitionConsumptionState.getOffsetRecord().getLeaderTopic(pubSubTopicRepository);
     int partitionId = partitionConsumptionState.getPartition();
-    if (partitionConsumptionState.getLeaderFollowerState().equals(LEADER) && leaderTopic != null) {
-      aggKafkaConsumerService
-          .unsubscribeConsumerFor(versionTopic, new PubSubTopicPartitionImpl(leaderTopic, partitionId));
-    } else {
+    PubSubTopic leaderTopic = partitionConsumptionState.getOffsetRecord().getLeaderTopic(pubSubTopicRepository);
+
+    // During leader/follower transitions or bootstrap a partition can be transiently subscribed
+    // to BOTH the version topic and a separate leader topic (e.g., RT). Unsubscribe each that is
+    // currently subscribed; consumerHasSubscription() avoids redundant work when a topic isn't
+    // attached.
+    if (consumerHasSubscription(versionTopic, partitionConsumptionState)) {
       aggKafkaConsumerService
           .unsubscribeConsumerFor(versionTopic, new PubSubTopicPartitionImpl(versionTopic, partitionId));
     }
+    if (leaderTopic != null && !leaderTopic.equals(versionTopic)
+        && consumerHasSubscription(leaderTopic, partitionConsumptionState)) {
+      aggKafkaConsumerService
+          .unsubscribeConsumerFor(versionTopic, new PubSubTopicPartitionImpl(leaderTopic, partitionId));
+      // Gate sep-RT unsubscribe on actual subscription to avoid WARN spam from the consumer
+      // delegator when sep-RT isn't currently attached (common during transitions).
+      if (isSeparatedRealtimeTopicEnabled() && leaderTopic.isRealTime() && separateRealTimeTopic != null
+          && consumerHasSubscription(separateRealTimeTopic, partitionConsumptionState)) {
+        aggKafkaConsumerService
+            .unsubscribeConsumerFor(versionTopic, new PubSubTopicPartitionImpl(separateRealTimeTopic, partitionId));
+      }
+    }
 
     /**
      * Leader of the user partition should close all subPartitions it is producing to.

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

@@ -216,7 +216,18 @@ enum LatchStatus {
    * and check timeout in the consumption state check function which runs regularly:
    * {@link LeaderFollowerStoreIngestionTask#checkLongRunningTaskState()}
    */
-  private final long consumptionStartTimeInMs;
+  private long consumptionStartTimeInMs;
+
+  /**
+   * Tracks whether this partition is currently paused due to a store-level
+   * {@link com.linkedin.venice.meta.IngestionPauseMode} other than NOT_PAUSED. Distinct from the
+   * disk-quota pause managed by {@link StorageUtilizationManager} — when this flag is true, quota
+   * resume callbacks must no-op so they don't fight the store-level pause.
+   * <p>
+   * Written by the SIT thread but read by disk-quota callbacks invoked from other threads, so it
+   * is {@code volatile} for cross-thread visibility.
+   */
+  private volatile boolean storeLevelPaused = false;
 
   /**
    * This hash map will keep a temporary mapping between a key and it's value.
@@ -956,6 +967,19 @@ public long getConsumptionStartTimeInMs() {
     return consumptionStartTimeInMs;
   }
 
+  /** Resets the bootstrap-timeout clock so paused time isn't counted against the window. */
+  public void resetConsumptionStartTimeInMs() {
+    this.consumptionStartTimeInMs = System.currentTimeMillis();
+  }
+
+  public boolean isStoreLevelPaused() {
+    return storeLevelPaused;
+  }
+
+  public void setStoreLevelPaused(boolean storeLevelPaused) {
+    this.storeLevelPaused = storeLevelPaused;
+  }
+
   public void setTransientRecord(
       int kafkaClusterId,
       PubSubPosition consumedPosition,

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

@@ -343,13 +343,19 @@ private boolean isStorageQuotaExceeded(StoragePartitionDiskUsage partitionDiskUs
    * partition without affecting partition subscription
    */
   private void pausePartition(int partition, String consumingTopic) {
-    pausePartition.execute(consumingTopic, partition);
-    this.pausedPartitions.add(partition);
+    // Only record the partition as quota-paused if the underlying consumer was actually paused.
+    // The lambda may no-op if a higher-priority pause source (e.g. store-level IngestionPauseMode)
+    // already owns the consumer; recording it as paused anyway would desynchronize this manager
+    // from real consumer state.
+    if (pausePartition.execute(consumingTopic, partition)) {
+      this.pausedPartitions.add(partition);
+    }
   }
 
   private void resumePartition(int partition, String consumingTopic) {
-    resumePartition.execute(consumingTopic, partition);
-    this.pausedPartitions.remove(partition);
+    if (resumePartition.execute(consumingTopic, partition)) {
+      this.pausedPartitions.remove(partition);
+    }
   }
 
   private void resumeAllPartitions() {

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

@@ -93,6 +93,7 @@
 import com.linkedin.venice.kafka.protocol.state.StoreVersionState;
 import com.linkedin.venice.message.KafkaKey;
 import com.linkedin.venice.meta.HybridStoreConfig;
+import com.linkedin.venice.meta.IngestionPauseMode;
 import com.linkedin.venice.meta.ReadOnlySchemaRepository;
 import com.linkedin.venice.meta.ReadOnlyStoreRepository;
 import com.linkedin.venice.meta.Store;
@@ -4671,16 +4672,39 @@ public void consumerResetOffset(PubSubTopic topic, PartitionConsumptionState par
     aggKafkaConsumerService.resetOffsetFor(versionTopic, new PubSubTopicPartitionImpl(topic, partitionId));
   }
 
-  private void pauseConsumption(String topic, int partitionId) {
+  private boolean pauseConsumption(String topic, int partitionId) {
+    // Store-level pause takes precedence; no-op here so the two pause sources don't race.
+    if (shouldSkipQuotaCallbackForStoreLevelPause(partitionConsumptionStateMap.get(partitionId))) {
+      logQuotaCallbackSuppressed("pauseConsumption", topic, partitionId);
+      return false;
+    }
     aggKafkaConsumerService.pauseConsumerFor(
         versionTopic,
         new PubSubTopicPartitionImpl(pubSubTopicRepository.getTopic(topic), partitionId));
+    return true;
   }
 
-  private void resumeConsumption(String topic, int partitionId) {
+  private boolean resumeConsumption(String topic, int partitionId) {
+    // Store-level pause takes precedence; no-op here so a quota resume doesn't un-pause us.
+    if (shouldSkipQuotaCallbackForStoreLevelPause(partitionConsumptionStateMap.get(partitionId))) {
+      logQuotaCallbackSuppressed("resumeConsumption", topic, partitionId);
+      return false;
+    }
     aggKafkaConsumerService.resumeConsumerFor(
         versionTopic,
         new PubSubTopicPartitionImpl(pubSubTopicRepository.getTopic(topic), partitionId));
+    return true;
+  }
+
+  private void logQuotaCallbackSuppressed(String callbackName, String topic, int partitionId) {
+    String key = storeName + "-" + callbackName + "-storeLevelPauseSuppressed";
+    if (!REDUNDANT_LOGGING_FILTER.isRedundantException(key)) {
+      LOGGER.info(
+          "Disk-quota {} suppressed for {}-{} because partition is store-level paused.",
+          callbackName,
+          topic,
+          partitionId);
+    }
   }
 
   /**
@@ -4753,6 +4777,42 @@ protected final void recordActiveKeyCountInvalidation() {
     hostLevelIngestionStats.recordActiveKeyCountInvalidation();
   }
 
+  /**
+   * Returns true when a disk-quota pause/resume callback must no-op because the partition is
+   * currently store-level paused. Pure helper to keep the branch unit-testable.
+   */
+  static boolean shouldSkipQuotaCallbackForStoreLevelPause(PartitionConsumptionState pcs) {
+    return pcs != null && pcs.isStoreLevelPaused();
+  }
+
+  /**
+   * Returns true if this SIT should pause based on the store's current pause mode and, for
+   * {@link IngestionPauseMode#CURRENT_VERSION}, whether this SIT's version number equals
+   * {@link Store#getCurrentVersion()}. Applies uniformly to both Venice servers and DaVinci
+   * clients.
+   */
+  boolean shouldPauseForStore(Store store) {
+    return shouldPauseForStore(store, versionNumber);
+  }
+
+  /**
+   * Static, side-effect-free pause decision. Exposed package-private for unit testing.
+   * @param store store metadata (must be non-null)
+   * @param sitVersionNumber the version of the SIT making the decision
+   * @return true if this SIT should pause Kafka consumption
+   */
+  static boolean shouldPauseForStore(Store store, int sitVersionNumber) {
+    IngestionPauseMode mode = store.getIngestionPauseMode();
+    if (mode == null || mode == IngestionPauseMode.NOT_PAUSED) {
+      return false;
+    }
+    if (mode == IngestionPauseMode.ALL_VERSIONS) {
+      return true;
+    }
+    // CURRENT_VERSION: pause only if this SIT's version is the current one
+    return sitVersionNumber == store.getCurrentVersion();
+  }
+
   /**
    * Write the kafka message to the underlying storage engine.
    * @param consumerRecord
@@ -6130,6 +6190,26 @@ && getResubscribeRequestQueue().peek() != null) {
         continue;
       }
 
+      /**
+       * Skip partitions whose store is currently paused. Resubscribing here would re-attach the
+       * Kafka consumer that the store-level pause has deliberately torn down; maybeTransitionPauseState
+       * would then re-unsubscribe it on the next reconcile tick, but the brief window allows
+       * records to leak through. The reconcile loop will resubscribe via {@link #resubscribe} on
+       * resume, so this lag-based path can safely defer.
+       */
+      if (pcs.isStoreLevelPaused()) {
+        // Throttle via the shared filter — during an operational pause every lagging partition
+        // gets enqueued each heartbeat cycle (~60s), so without throttling a paused store with
+        // N partitions emits N log lines per cycle. Key on storeName so we get one signal per
+        // paused store per filter window instead of one per replica.
+        if (!REDUNDANT_LOGGING_FILTER.isRedundantException(storeName + "-skip-lag-resubscribe-store-paused")) {
+          LOGGER.info(
+              "Skipping lag-based resubscribe for replica: {} because store-level pause is active.",
+              pcs.getReplicaId());
+        }
+        continue;
+      }
+
       /**
        * Skip partitions with an in-flight blob transfer. Resubscribing reopens RocksDB on the final
        * partition directory while blob transfer is still receiving files into the temp directory,