[dvc][cc] Fix unsafe DVRT CDC version swap

.github/workflows/VeniceCI-E2ETests.yml

@@ -1629,6 +1629,44 @@ jobs:
         with:
           shard: 85
 
+  IntegrationTests_86:
+    name: IntegrationTests_86
+    strategy:
+      fail-fast: false
+    runs-on: ubuntu-latest
+    permissions:
+     contents: read
+    timeout-minutes: 15
+    concurrency:
+     group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-jdk17-IntegrationTests_86
+     cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+      - uses: ./.github/actions/integration-test-run
+        with:
+          shard: 86
+
+  IntegrationTests_87:
+    name: IntegrationTests_87
+    strategy:
+      fail-fast: false
+    runs-on: ubuntu-latest
+    permissions:
+     contents: read
+    timeout-minutes: 15
+    concurrency:
+     group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-jdk17-IntegrationTests_87
+     cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+      - uses: ./.github/actions/integration-test-run
+        with:
+          shard: 87
+
   IntegrationTests_99:
     name: IntegrationTests_99
     strategy:
@@ -1656,7 +1694,7 @@ jobs:
     runs-on: ubuntu-latest
     permissions:
       contents: read
-    needs: [IntegrationTests_1, IntegrationTests_2, IntegrationTests_3, IntegrationTests_4, IntegrationTests_5, IntegrationTests_6, IntegrationTests_7, IntegrationTests_8, IntegrationTests_9, IntegrationTests_10, IntegrationTests_11, IntegrationTests_12, IntegrationTests_13, IntegrationTests_14, IntegrationTests_15, IntegrationTests_16, IntegrationTests_17, IntegrationTests_18, IntegrationTests_19, IntegrationTests_20, IntegrationTests_21, IntegrationTests_22, IntegrationTests_23, IntegrationTests_24, IntegrationTests_25, IntegrationTests_26, IntegrationTests_27, IntegrationTests_28, IntegrationTests_29, IntegrationTests_30, IntegrationTests_31, IntegrationTests_32, IntegrationTests_33, IntegrationTests_34, IntegrationTests_35, IntegrationTests_36, IntegrationTests_37, IntegrationTests_38, IntegrationTests_39, IntegrationTests_40, IntegrationTests_41, IntegrationTests_42, IntegrationTests_43, IntegrationTests_44, IntegrationTests_45, IntegrationTests_46, IntegrationTests_47, IntegrationTests_48, IntegrationTests_49, IntegrationTests_50, IntegrationTests_51, IntegrationTests_52, IntegrationTests_53, IntegrationTests_54, IntegrationTests_55, IntegrationTests_56, IntegrationTests_57, IntegrationTests_58, IntegrationTests_59, IntegrationTests_60, IntegrationTests_61, IntegrationTests_62, IntegrationTests_63, IntegrationTests_64, IntegrationTests_65, IntegrationTests_66, IntegrationTests_67, IntegrationTests_68, IntegrationTests_69, IntegrationTests_70, IntegrationTests_71, IntegrationTests_72, IntegrationTests_73, IntegrationTests_74, IntegrationTests_75, IntegrationTests_76, IntegrationTests_77, IntegrationTests_78, IntegrationTests_79, IntegrationTests_80, IntegrationTests_81, IntegrationTests_82, IntegrationTests_83, IntegrationTests_84, IntegrationTests_85, IntegrationTests_99]
+    needs: [IntegrationTests_1, IntegrationTests_2, IntegrationTests_3, IntegrationTests_4, IntegrationTests_5, IntegrationTests_6, IntegrationTests_7, IntegrationTests_8, IntegrationTests_9, IntegrationTests_10, IntegrationTests_11, IntegrationTests_12, IntegrationTests_13, IntegrationTests_14, IntegrationTests_15, IntegrationTests_16, IntegrationTests_17, IntegrationTests_18, IntegrationTests_19, IntegrationTests_20, IntegrationTests_21, IntegrationTests_22, IntegrationTests_23, IntegrationTests_24, IntegrationTests_25, IntegrationTests_26, IntegrationTests_27, IntegrationTests_28, IntegrationTests_29, IntegrationTests_30, IntegrationTests_31, IntegrationTests_32, IntegrationTests_33, IntegrationTests_34, IntegrationTests_35, IntegrationTests_36, IntegrationTests_37, IntegrationTests_38, IntegrationTests_39, IntegrationTests_40, IntegrationTests_41, IntegrationTests_42, IntegrationTests_43, IntegrationTests_44, IntegrationTests_45, IntegrationTests_46, IntegrationTests_47, IntegrationTests_48, IntegrationTests_49, IntegrationTests_50, IntegrationTests_51, IntegrationTests_52, IntegrationTests_53, IntegrationTests_54, IntegrationTests_55, IntegrationTests_56, IntegrationTests_57, IntegrationTests_58, IntegrationTests_59, IntegrationTests_60, IntegrationTests_61, IntegrationTests_62, IntegrationTests_63, IntegrationTests_64, IntegrationTests_65, IntegrationTests_66, IntegrationTests_67, IntegrationTests_68, IntegrationTests_69, IntegrationTests_70, IntegrationTests_71, IntegrationTests_72, IntegrationTests_73, IntegrationTests_74, IntegrationTests_75, IntegrationTests_76, IntegrationTests_77, IntegrationTests_78, IntegrationTests_79, IntegrationTests_80, IntegrationTests_81, IntegrationTests_82, IntegrationTests_83, IntegrationTests_84, IntegrationTests_85, IntegrationTests_86, IntegrationTests_87, IntegrationTests_99]
     timeout-minutes: 20
     if: ${{ cancelled() || contains(needs.*.result, 'cancelled') || contains(needs.*.result, 'failure') }}
     steps:

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

@@ -5,6 +5,7 @@
 import com.linkedin.venice.annotation.Experimental;
 import com.linkedin.venice.compression.VeniceCompressor;
 import com.linkedin.venice.kafka.protocol.ControlMessage;
+import com.linkedin.venice.kafka.protocol.VersionSwap;
 import com.linkedin.venice.kafka.protocol.state.PartitionState;
 import com.linkedin.venice.meta.ReadOnlySchemaRepository;
 import com.linkedin.venice.pubsub.PubSubContext;
@@ -14,6 +15,7 @@
 import java.io.IOException;
 import java.util.Map;
 import java.util.concurrent.CountDownLatch;
+import java.util.function.IntConsumer;
 import org.apache.avro.Schema;
 
 
@@ -34,6 +36,13 @@ public class InternalDaVinciRecordTransformer<K, V, O> extends DaVinciRecordTran
    */
   private final CountDownLatch startLatchConsumptionLatch;
 
+  /**
+   * Wired by {@link com.linkedin.davinci.kafka.consumer.StoreIngestionTask} so that the CDC consumer
+   * can pause/resume Kafka prefetch on its own version's partitions during the AA version swap barrier.
+   */
+  private IntConsumer pausePartitionConsumer;
+  private IntConsumer resumePartitionConsumer;
+
   public InternalDaVinciRecordTransformer(
       DaVinciRecordTransformer recordTransformer,
       Schema keySchema,
@@ -52,6 +61,20 @@ public InternalDaVinciRecordTransformer(
         new CountDownLatch(internalRecordTransformerConfig.getStartConsumptionLatchCount());
   }
 
+  /**
+   * Wires the back-reference from the CDC inner-class transformer to this {@code InternalDaVinciRecordTransformer},
+   * enabling the AA version-swap coordinator to drive Kafka prefetch pause/resume on this version's partitions.
+   * Called from {@link com.linkedin.davinci.kafka.consumer.StoreIngestionTask} after construction; calling this
+   * after construction (rather than inside it) avoids invoking an overridable method on the user's transformer
+   * before initialization is complete.
+   */
+  public void initializeChangelogConsumerBackReference() {
+    if (isCDCRecordTransformer()) {
+      ((VeniceChangelogConsumerDaVinciRecordTransformerImpl.DaVinciRecordTransformerChangelogConsumer) this.recordTransformer)
+          .setInternalRecordTransformer(this);
+    }
+  }
+
   @Override
   public DaVinciRecordTransformerResult<O> transform(
       Lazy<K> key,
@@ -91,13 +114,44 @@ public boolean useUniformInputValueSchema() {
   }
 
   /**
-   * Lifecycle event triggered when a version swap is detected for partitionId
-   * It is used for DVRT CDC.
+   * Lifecycle event triggered when a version swap is detected for partitionId.
+   * It is used for the DaVinciRecordTransformer CDC consumer. The {@code versionSwap} payload is
+   * required for AA-aware version-swap coordination (region filtering, generation-id matching);
+   * it may be null when invoked from legacy or test code paths.
    */
-  public void onVersionSwap(int currentVersion, int futureVersion, int partitionId) {
+  public void onVersionSwap(VersionSwap versionSwap, int currentVersion, int futureVersion, int partitionId) {
     if (isCDCRecordTransformer()) {
       ((VeniceChangelogConsumerDaVinciRecordTransformerImpl.DaVinciRecordTransformerChangelogConsumer) this.recordTransformer)
-          .onVersionSwap(currentVersion, futureVersion, partitionId);
+          .onVersionSwap(versionSwap, currentVersion, futureVersion, partitionId);
+    }
+  }
+
+  /**
+   * Wires the pause/resume handlers used by the AA version-swap coordinator to drive Kafka prefetch
+   * pause/resume on this transformer's own version's partitions.
+   */
+  public void setPartitionPauseHandlers(IntConsumer pauseHandler, IntConsumer resumeHandler) {
+    this.pausePartitionConsumer = pauseHandler;
+    this.resumePartitionConsumer = resumeHandler;
+  }
+
+  /**
+   * Pauses Kafka prefetch for {@code partitionId} on this transformer's version. No-op when handlers
+   * are not wired (e.g. unit tests or non-server contexts).
+   */
+  public void pausePartitionConsumption(int partitionId) {
+    if (pausePartitionConsumer != null) {
+      pausePartitionConsumer.accept(partitionId);
+    }
+  }
+
+  /**
+   * Resumes Kafka prefetch for {@code partitionId} on this transformer's version. No-op when handlers
+   * are not wired.
+   */
+  public void resumePartitionConsumption(int partitionId) {
+    if (resumePartitionConsumer != null) {
+      resumePartitionConsumer.accept(partitionId);
     }
   }