[test] [server] Fix Flaky Test testProcessConsumerActionsError() (#1527)

Fix the flaky test `testProcessConsumerActionsError()` by splitting the error handling off into a dedicated method and unit testing that directly. 💥

Author: KaiSernLim

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

@@ -1960,42 +1960,46 @@ void processConsumerActions(Store store) throws InterruptedException {
         action.getFuture().completeExceptionally(e);
         throw e;
       } catch (Throwable e) {
-        if (action.getAttemptsCount() <= MAX_CONSUMER_ACTION_ATTEMPTS) {
-          LOGGER.warn("Failed to process consumer action {}, will retry later.", action, e);
+        if (!handleConsumerActionsError(e, action, actionProcessStartTimeInMs)) {
           return;
         }
-        LOGGER.error(
-            "Failed to execute consumer action {} after {} attempts. Total elapsed time: {}ms",
-            action,
-            action.getAttemptsCount(),
-            LatencyUtils.getElapsedTimeFromMsToMs(actionProcessStartTimeInMs),
-            e);
-        // Mark action as failed since it has exhausted all the retries.
-        action.getFuture().completeExceptionally(e);
-        // After MAX_CONSUMER_ACTION_ATTEMPTS retries we should give up and error the ingestion task.
-        PartitionConsumptionState state = partitionConsumptionStateMap.get(action.getPartition());
-
-        // Remove the action that is failed to execute recently (not necessarily the head of consumerActionsQueue).
-        if (consumerActionsQueue.remove(action)) {
-          partitionToPendingConsumerActionCountMap.get(action.getPartition()).decrementAndGet();
-        }
-        /**
-         * {@link state} can be null if the {@link OffsetRecord} from {@link storageMetadataService} was corrupted in
-         * {@link #processCommonConsumerAction}, so the {@link PartitionConsumptionState} was never created
-         */
-        if (state == null || !state.isCompletionReported()) {
-          reportError(
-              "Error when processing consumer action: " + action,
-              action.getPartition(),
-              new VeniceException(e));
-        }
       }
     }
     if (emitMetrics.get()) {
       hostLevelIngestionStats.recordProcessConsumerActionLatency(Duration.between(startTime, Instant.now()).toMillis());
     }
   }
 
+  boolean handleConsumerActionsError(Throwable e, ConsumerAction action, long actionProcessStartTimeInMs) {
+    if (action.getAttemptsCount() <= MAX_CONSUMER_ACTION_ATTEMPTS) {
+      LOGGER.warn("Failed to process consumer action {}, will retry later.", action, e);
+      return false;
+    }
+    LOGGER.error(
+        "Failed to execute consumer action {} after {} attempts. Total elapsed time: {}ms",
+        action,
+        action.getAttemptsCount(),
+        LatencyUtils.getElapsedTimeFromMsToMs(actionProcessStartTimeInMs),
+        e);
+    // Mark action as failed since it has exhausted all the retries.
+    action.getFuture().completeExceptionally(e);
+    // After MAX_CONSUMER_ACTION_ATTEMPTS retries we should give up and error the ingestion task.
+    PartitionConsumptionState state = partitionConsumptionStateMap.get(action.getPartition());
+
+    // Remove the action that is failed to execute recently (not necessarily the head of consumerActionsQueue).
+    if (consumerActionsQueue.remove(action)) {
+      partitionToPendingConsumerActionCountMap.get(action.getPartition()).decrementAndGet();
+    }
+    /**
+     * {@link state} can be null if the {@link OffsetRecord} from {@link storageMetadataService} was corrupted in
+     * {@link #processCommonConsumerAction}, so the {@link PartitionConsumptionState} was never created
+     */
+    if (state == null || !state.isCompletionReported()) {
+      reportError("Error when processing consumer action: " + action, action.getPartition(), new VeniceException(e));
+    }
+    return true;
+  }
+
   /**
    * Applies name resolution to all Kafka URLs in the provided TopicSwitch. Useful for translating URLs that came from
    * a different runtime (e.g. from the controller, or from state persisted by a previous run of the same server).

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

@@ -5321,31 +5321,27 @@ public void testMeasureLagWithCallToPubSub() {
   /**
    * When SIT encounters a corrupted {@link OffsetRecord} in {@link StoreIngestionTask#processCommonConsumerAction} and
    * {@link StorageMetadataService#getLastOffset} throws an exception due to a deserialization error,
-   * {@link StoreIngestionTask#reportError(String, int, Exception)} should be called in order to trigger a Helix
-   * state transition without waiting 24+ hours for the Helix state transition timeout.
+   * {@link StoreIngestionTask#reportError(String, int, Exception)} from
+   * {@link StoreIngestionTask#handleConsumerActionsError(Throwable, ConsumerAction, long)} should be called in
+   * order to trigger a Helix state transition without waiting 24+ hours for the Helix state transition timeout.
    */
   @Test
-  public void testProcessConsumerActionsError() throws Exception {
-    runTest(Collections.singleton(PARTITION_FOO), () -> {
-      storeIngestionTaskUnderTest.close(); // prevent the SIT polling thread run from interfering with the
-                                           // processConsumerActions()
-
+  public void testHandleConsumerActionsError() throws Exception {
+    final int p = 0; // partition number
+    runTest(Collections.singleton(p), () -> {
       // This is an actual exception thrown when deserializing a corrupted OffsetRecord
       String msg = "Received Magic Byte '6' which is not supported by InternalAvroSpecificSerializer. "
           + "The only supported Magic Byte for this implementation is '24'.";
-      when(mockStorageMetadataService.getLastOffset(any(), anyInt())).thenThrow(new VeniceMessageException(msg));
-      // To reach reportError(), bypass the conditional: action.getAttemptsCount() <= MAX_CONSUMER_ACTION_ATTEMPTS
-      for (int i = 0; i < StoreIngestionTask.MAX_CONSUMER_ACTION_ATTEMPTS + 1; i++) {
-        try {
-          storeIngestionTaskUnderTest.processConsumerActions(storeAndVersionConfigsUnderTest.store);
-        } catch (InterruptedException e) {
-          throw new RuntimeException(e);
-        }
-      }
-      ArgumentCaptor<VeniceException> captor = ArgumentCaptor.forClass(VeniceException.class);
-      verify(storeIngestionTaskUnderTest, timeout(TEST_TIMEOUT_MS).atLeast(1))
-          .reportError(anyString(), eq(PARTITION_FOO), captor.capture());
-      assertTrue(captor.getValue().getMessage().endsWith(msg));
+      Throwable e = new VeniceException(msg);
+      CompletableFuture<Void> future = new CompletableFuture<>();// CompletableFuture(null);
+      ConsumerAction action = mock(ConsumerAction.class);
+      when(action.getPartition()).thenReturn(p);
+      when(action.getFuture()).thenReturn(future);
+      when(action.getAttemptsCount()).thenReturn(StoreIngestionTask.MAX_CONSUMER_ACTION_ATTEMPTS + 1);
+      assertFalse(future.isCompletedExceptionally());
+      storeIngestionTaskUnderTest.handleConsumerActionsError(e, action, 0L);
+      verify(storeIngestionTaskUnderTest, times(1)).reportError(anyString(), eq(p), any(VeniceException.class));
+      assertTrue(future.isCompletedExceptionally());
     }, AA_OFF);
   }