add test for isDavinciHeartbeatReported

linkedin · Dec 11, 2024 · 80adb8b · 80adb8b
1 parent f84707f
commit 80adb8b
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 2 deletions.
diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/meta/AbstractStore.java b/internal/venice-common/src/main/java/com/linkedin/venice/meta/AbstractStore.java
@@ -351,9 +351,13 @@ public void fixMissingFields() {
 
   @Override
   public void updateVersionForDaVinciHeartbeat(int versionNumber, boolean reported) {
-    for (StoreVersion version: storeVersionsSupplier.getForUpdate()) {
+    checkVersionSupplier();
+
+    for (StoreVersion storeVersion: storeVersionsSupplier.getForUpdate()) {
+      Version version = new VersionImpl(storeVersion);
       if (version.getNumber() == versionNumber) {
-        version.setIsDaVinciHeartBeatReported(reported);
+        version.setIsDavinciHeartbeatReported(reported);
+        return;
       }
     }
   }

diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/meta/VersionImpl.java b/internal/venice-common/src/main/java/com/linkedin/venice/meta/VersionImpl.java
@@ -507,6 +507,7 @@ public Version cloneVersion() {
     clonedVersion.setBlobTransferEnabled(isBlobTransferEnabled());
     clonedVersion.setTargetSwapRegion(getTargetSwapRegion());
     clonedVersion.setTargetSwapRegionWaitTime(getTargetSwapRegionWaitTime());
+    clonedVersion.setIsDavinciHeartbeatReported(getIsDavinciHeartbeatReported());
     return clonedVersion;
   }
 

diff --git a/...-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientTest.java b/...-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientTest.java
@@ -75,6 +75,7 @@
 import com.linkedin.venice.integration.utils.VeniceRouterWrapper;
 import com.linkedin.venice.meta.IngestionMetadataUpdateType;
 import com.linkedin.venice.meta.IngestionMode;
+import com.linkedin.venice.meta.StoreInfo;
 import com.linkedin.venice.meta.Version;
 import com.linkedin.venice.offsets.OffsetRecord;
 import com.linkedin.venice.partitioner.ConstantVenicePartitioner;
@@ -1547,6 +1548,32 @@ public void testDVCSnapshotGeneration(boolean useDiskStorage) throws Exception {
     }
   }
 
+  @Test
+  public void testIsDavinciHeartbeatReported() throws Exception {
+    String storeName = Utils.getUniqueString("testIsDavinviHeartbeatReported");
+    Consumer<UpdateStoreQueryParams> paramsConsumer = params -> params.setTargetRegionSwap("test");
+    setUpStore(storeName, paramsConsumer, properties -> {}, true);
+
+    try (DaVinciClient<Object, Object> client = ServiceFactory.getGenericAvroDaVinciClient(storeName, cluster)) {
+      client.subscribeAll().get();
+    }
+
+    Properties vpjProperties = defaultVPJProps(cluster, inputDirPath, storeName);
+    runVPJ(vpjProperties, 2, cluster);
+
+    try (ControllerClient controllerClient = cluster.getControllerClient()) {
+      StoreInfo store = controllerClient.getStore(storeName).getStore();
+      TestUtils.waitForNonDeterministicAssertion(
+          2,
+          TimeUnit.MILLISECONDS,
+          () -> Assert.assertEquals(store.getIsDavinciHeartbeatReported(), true));
+      TestUtils.waitForNonDeterministicAssertion(
+          2,
+          TimeUnit.MILLISECONDS,
+          () -> Assert.assertEquals(store.getVersion(2).get().getIsDavinciHeartbeatReported(), true));
+    }
+  }
+
   private void setupHybridStore(String storeName, Consumer<UpdateStoreQueryParams> paramsConsumer) throws Exception {
     setupHybridStore(storeName, paramsConsumer, KEY_COUNT);
   }

diff --git a/.../venice-controller/src/main/java/com/linkedin/venice/pushmonitor/PushStatusCollector.java b/.../venice-controller/src/main/java/com/linkedin/venice/pushmonitor/PushStatusCollector.java
@@ -178,6 +178,7 @@ private void scanDaVinciPushStatus() {
       if (!store.getIsDavinciHeartbeatReported() && daVinciStatus.getStatus() != ExecutionStatus.NOT_CREATED) {
         int versionNum = Version.parseVersionFromVersionTopicName(pushStatus.topicName);
         store.updateVersionForDaVinciHeartbeat(versionNum, true);
+        store.setIsDavinciHeartbeatReported(true);
         storeRepository.updateStore(store);
       }