add unit test

Author: sixpluszero

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

@@ -5246,12 +5246,13 @@ && getResubscribeRequestQueue().peek() != null) {
         continue;
       }
       try {
+        LOGGER.info("Resubscribing: {}", pcs.getReplicaId());
         getPartitionToPreviousResubscribeTimeMap().put(partition, System.currentTimeMillis());
+        count++;
         resubscribe(pcs);
       } catch (Exception e) {
         LOGGER.warn("Caught exception when resubscribing for replica: {}", pcs.getReplicaId());
       }
-      count++;
     }
   }
 

clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/ingestion/heartbeat/HeartbeatMonitoringService.java

@@ -617,14 +617,14 @@ protected void record() {
             .recordFollowerLag(storeName, version, region, heartbeatTs, isReadyToServe)));
   }
 
-  protected void checkAndMaybeLogHeartbeatDelayMap(
+  void checkAndMaybeLogHeartbeatDelayMap(
       Map<String, Map<Integer, Map<Integer, Map<String, HeartbeatTimeStampEntry>>>> heartbeatTimestamps) {
-    if (kafkaStoreIngestionService == null) {
+    if (getKafkaStoreIngestionService() == null) {
       // Service not initialized yet, skip logging
       return;
     }
     long currentTimestamp = System.currentTimeMillis();
-    boolean isLeader = heartbeatTimestamps == leaderHeartbeatTimeStamps;
+    boolean isLeader = heartbeatTimestamps == getLeaderHeartbeatTimeStamps();
     for (Map.Entry<String, Map<Integer, Map<Integer, Map<String, HeartbeatTimeStampEntry>>>> storeName: heartbeatTimestamps
         .entrySet()) {
       for (Map.Entry<Integer, Map<Integer, Map<String, HeartbeatTimeStampEntry>>> version: storeName.getValue()
@@ -644,14 +644,18 @@ protected void checkAndMaybeLogHeartbeatDelayMap(
                   lag,
                   heartbeatTs,
                   currentTimestamp);
-              kafkaStoreIngestionService.attemptToPrintIngestionInfoFor(
+              getKafkaStoreIngestionService().attemptToPrintIngestionInfoFor(
                   storeName.getKey(),
                   version.getKey(),
                   partition.getKey(),
                   region.getKey());
-              if (serverConfig.isLagBasedReplicaAutoResubscribeEnabled()
-                  && serverConfig.getLagBasedReplicaAutoResubscribeThresholdInSeconds() < lag) {
-                kafkaStoreIngestionService
+              /**
+               * Here we don't consider whether it is current version or not, as it will need extra logic to extract.
+               * We will delegate to KafkaConsumerService to determine whether it takes this request as it has information.
+               */
+              if (getServerConfig().isLagBasedReplicaAutoResubscribeEnabled() && TimeUnit.SECONDS
+                  .toMillis(getServerConfig().getLagBasedReplicaAutoResubscribeThresholdInSeconds()) < lag) {
+                getKafkaStoreIngestionService()
                     .maybeAddResubscribeRequest(storeName.getKey(), version.getKey(), partition.getKey());
               }
             }
@@ -894,4 +898,12 @@ String getLocalRegionName() {
   public void setKafkaStoreIngestionService(KafkaStoreIngestionService kafkaStoreIngestionService) {
     this.kafkaStoreIngestionService = kafkaStoreIngestionService;
   }
+
+  KafkaStoreIngestionService getKafkaStoreIngestionService() {
+    return kafkaStoreIngestionService;
+  }
+
+  VeniceServerConfig getServerConfig() {
+    return serverConfig;
+  }
 }

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

@@ -0,0 +1,58 @@
+package com.linkedin.davinci.kafka.consumer;
+
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doCallRealMethod;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import com.linkedin.davinci.config.VeniceServerConfig;
+import com.linkedin.venice.utils.concurrent.VeniceConcurrentHashMap;
+import java.util.Map;
+import java.util.concurrent.PriorityBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import org.testng.annotations.Test;
+
+
+public class AutoResubscribeTest {
+  @Test
+  public void testHandleAutoResubscribe() throws InterruptedException {
+    StoreIngestionTask storeIngestionTask = mock(StoreIngestionTask.class);
+    doCallRealMethod().when(storeIngestionTask).maybeProcessResubscribeRequest();
+    VeniceServerConfig serverConfig = mock(VeniceServerConfig.class);
+    doReturn(serverConfig).when(storeIngestionTask).getServerConfig();
+    PriorityBlockingQueue<Integer> resubscribeQueue = new PriorityBlockingQueue<>();
+    doReturn(resubscribeQueue).when(storeIngestionTask).getResubscribeRequestQueue();
+    Map<Integer, Long> resubscribeRequestTimestamp = new VeniceConcurrentHashMap<>();
+    doReturn(resubscribeRequestTimestamp).when(storeIngestionTask).getPartitionToPreviousResubscribeTimeMap();
+    Map<Integer, PartitionConsumptionState> partitionConsumptionStateMap = new VeniceConcurrentHashMap<>();
+    doReturn(partitionConsumptionStateMap).when(storeIngestionTask).getPartitionConsumptionStateMap();
+    doReturn(300).when(serverConfig).getLagBasedReplicaAutoResubscribeIntervalInSeconds();
+
+    doReturn(2).when(serverConfig).getLagBasedReplicaAutoResubscribeMaxReplicaCount();
+
+    PartitionConsumptionState pcs1 = mock(PartitionConsumptionState.class);
+    PartitionConsumptionState pcs3 = mock(PartitionConsumptionState.class);
+    PartitionConsumptionState pcs4 = mock(PartitionConsumptionState.class);
+    partitionConsumptionStateMap.put(1, pcs1);
+    partitionConsumptionStateMap.put(3, pcs3);
+    partitionConsumptionStateMap.put(4, pcs4);
+    resubscribeRequestTimestamp.put(0, System.currentTimeMillis());
+    resubscribeRequestTimestamp.put(1, System.currentTimeMillis() - TimeUnit.SECONDS.toMillis(500));
+    resubscribeRequestTimestamp.put(2, System.currentTimeMillis() - TimeUnit.SECONDS.toMillis(500));
+    resubscribeRequestTimestamp.put(4, System.currentTimeMillis() - TimeUnit.SECONDS.toMillis(500));
+    resubscribeQueue.put(0);
+    resubscribeQueue.put(1);
+    resubscribeQueue.put(2);
+    resubscribeQueue.put(3);
+    resubscribeQueue.put(4);
+
+    storeIngestionTask.maybeProcessResubscribeRequest();
+    verify(storeIngestionTask, times(1)).resubscribe(eq(pcs1));
+    verify(storeIngestionTask, times(1)).resubscribe(eq(pcs3));
+    verify(storeIngestionTask, never()).resubscribe(eq(pcs4));
+
+  }
+}

clients/da-vinci-client/src/test/java/com/linkedin/davinci/stats/ingestion/heartbeat/HeartbeatMonitoringServiceTest.java

@@ -22,6 +22,7 @@
 import static org.mockito.Mockito.when;
 
 import com.linkedin.davinci.config.VeniceServerConfig;
+import com.linkedin.davinci.kafka.consumer.KafkaStoreIngestionService;
 import com.linkedin.davinci.kafka.consumer.LeaderFollowerStateType;
 import com.linkedin.davinci.kafka.consumer.PartitionConsumptionState;
 import com.linkedin.venice.exceptions.VeniceNoHelixResourceException;
@@ -42,6 +43,7 @@
 import com.linkedin.venice.utils.concurrent.VeniceConcurrentHashMap;
 import io.tehuti.metrics.MetricsRepository;
 import java.time.Duration;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
@@ -783,4 +785,36 @@ public void testLargestHeartbeatLag() {
     Assert.assertEquals(aggregatedHeartbeatLagEntry.getCurrentVersionHeartbeatLag(), 8000L);
     Assert.assertEquals(aggregatedHeartbeatLagEntry.getNonCurrentVersionHeartbeatLag(), 9900L);
   }
+
+  @Test
+  public void testTriggerAutoResubscribe() {
+    String store = "foo";
+    int version = 100;
+    int partition = 123;
+    String region = "dc1";
+    Map<String, Map<Integer, Map<Integer, Map<String, HeartbeatTimeStampEntry>>>> heartbeatTimestamps = new HashMap<>();
+    HeartbeatTimeStampEntry entry =
+        new HeartbeatTimeStampEntry(System.currentTimeMillis() - TimeUnit.MINUTES.toMillis(15), true, true);
+    heartbeatTimestamps.put(store, new HashMap<>());
+    heartbeatTimestamps.get(store).put(version, new HashMap<>());
+    heartbeatTimestamps.get(store).get(version).put(partition, new HashMap<>());
+    heartbeatTimestamps.get(store).get(version).get(partition).put(region, entry);
+
+    HeartbeatMonitoringService heartbeatMonitoringService = mock(HeartbeatMonitoringService.class);
+    KafkaStoreIngestionService kafkaStoreIngestionService = mock(KafkaStoreIngestionService.class);
+    VeniceServerConfig serverConfig = mock(VeniceServerConfig.class);
+    doReturn(serverConfig).when(heartbeatMonitoringService).getServerConfig();
+    doReturn(kafkaStoreIngestionService).when(heartbeatMonitoringService).getKafkaStoreIngestionService();
+    doCallRealMethod().when(heartbeatMonitoringService).checkAndMaybeLogHeartbeatDelayMap(anyMap());
+
+    // Config not enabled, nothing happen
+    heartbeatMonitoringService.checkAndMaybeLogHeartbeatDelayMap(heartbeatTimestamps);
+    verify(kafkaStoreIngestionService, never()).maybeAddResubscribeRequest(eq(store), eq(version), eq(partition));
+
+    // Config enabled, trigger resubscribe.
+    doReturn(true).when(serverConfig).isLagBasedReplicaAutoResubscribeEnabled();
+    doReturn(600).when(serverConfig).getLagBasedReplicaAutoResubscribeThresholdInSeconds();
+    heartbeatMonitoringService.checkAndMaybeLogHeartbeatDelayMap(heartbeatTimestamps);
+    verify(kafkaStoreIngestionService, times(1)).maybeAddResubscribeRequest(eq(store), eq(version), eq(partition));
+  }
 }