[controller] Handle NPE for incremental push reporting (#677)

* throw exception with partition id when hit NPE
linkedin · Oct 3, 2023 · 97df1e7 · 97df1e7
1 parent 2396100
commit 97df1e7
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/pushmonitor/OfflinePushStatus.java b/internal/venice-common/src/main/java/com/linkedin/venice/pushmonitor/OfflinePushStatus.java
@@ -208,6 +208,10 @@ public Map<Integer, Map<CharSequence, Integer>> getIncrementalPushStatus(
     for (PartitionStatus partitionStatus: getPartitionStatuses()) {
       Map<CharSequence, Integer> partitionPushStatus = new HashMap<>();
       Partition partition = partitionAssignment.getPartition(partitionStatus.getPartitionId());
+      if (partition == null) {
+        throw new VeniceException(
+            "Partition " + partitionStatus.getPartitionId() + " not found in partition assignment");
+      }
       Set<String> workingInstances =
           partition.getWorkingInstances().stream().map(Instance::getNodeId).collect(Collectors.toSet());
       for (ReplicaStatus replicaStatus: partitionStatus.getReplicaStatuses()) {

diff --git a/...al/venice-common/src/test/java/com/linkedin/venice/pushmonitor/OfflinePushStatusTest.java b/...al/venice-common/src/test/java/com/linkedin/venice/pushmonitor/OfflinePushStatusTest.java
@@ -9,6 +9,7 @@
 
 import com.linkedin.venice.exceptions.VeniceException;
 import com.linkedin.venice.meta.OfflinePushStrategy;
+import com.linkedin.venice.meta.PartitionAssignment;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -163,6 +164,14 @@ public void testCloneOfflinePushStatus() {
     Assert.assertNotEquals(clonedPush, offlinePushStatus);
   }
 
+  @Test(expectedExceptions = VeniceException.class, expectedExceptionsMessageRegExp = "Partition 0 not found in partition assignment")
+  public void testNPECaughtWhenPollingIncPushStatus() {
+    OfflinePushStatus offlinePushStatus =
+        new OfflinePushStatus(kafkaTopic, numberOfPartition, replicationFactor, strategy);
+    PartitionAssignment partitionAssignment = new PartitionAssignment("test-topic", 10);
+    offlinePushStatus.getIncrementalPushStatus(partitionAssignment, "ignore");
+  }
+
   private void testValidTargetStatuses(ExecutionStatus from, ExecutionStatus... statuses) {
     for (ExecutionStatus status: statuses) {
       OfflinePushStatus offlinePushStatus =