Handle nulls in DruidCoordinator.getReplicationFactor (#14447)

apache · Jun 20, 2023 · f5cc823 · f5cc823
1 parent 09d6c5a
commit f5cc823
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/server/src/main/java/org/apache/druid/server/coordinator/DruidCoordinator.java b/server/src/main/java/org/apache/druid/server/coordinator/DruidCoordinator.java
@@ -302,11 +302,11 @@ public Map<String, Double> getDatasourceToLoadStatus()
   @Nullable
   public Integer getReplicationFactor(SegmentId segmentId)
   {
-    if (segmentReplicationStatus != null) {
-      return segmentReplicationStatus.getReplicaCountsInCluster(segmentId).required();
-    } else {
+    if (segmentReplicationStatus == null) {
       return null;
     }
+    SegmentReplicaCount replicaCountsInCluster = segmentReplicationStatus.getReplicaCountsInCluster(segmentId);
+    return replicaCountsInCluster == null ? null : replicaCountsInCluster.required();
   }
 
   @Nullable

diff --git a/server/src/test/java/org/apache/druid/server/coordinator/DruidCoordinatorTest.java b/server/src/test/java/org/apache/druid/server/coordinator/DruidCoordinatorTest.java
@@ -249,6 +249,9 @@ public void testCoordinatorRun() throws Exception
     EasyMock.replay(serverInventoryView, loadQueueTaskMaster);
 
     coordinator.start();
+
+    Assert.assertNull(coordinator.getReplicationFactor(dataSegment.getId()));
+
     // Wait for this coordinator to become leader
     leaderAnnouncerLatch.await();
 
@@ -306,6 +309,7 @@ public void testCoordinatorRun() throws Exception
     // the segments are replicated as many times as they can be given state of cluster, therefore should not be
     // under-replicated.
     Assert.assertEquals(0L, underRepliicationCountsPerDataSourceUsingClusterView.getLong(dataSource));
+    Assert.assertEquals(Integer.valueOf(2), coordinator.getReplicationFactor(dataSegment.getId()));
 
     coordinator.stop();
     leaderUnannouncerLatch.await();
@@ -398,6 +402,8 @@ public void testCoordinatorTieredRun() throws Exception
     Assert.assertEquals(0L, underReplicationCountsPerDataSourcePerTierUsingClusterView.get(hotTierName).getLong(dataSource));
     Assert.assertEquals(0L, underReplicationCountsPerDataSourcePerTierUsingClusterView.get(coldTierName).getLong(dataSource));
 
+    dataSegments.values().forEach(dataSegment -> Assert.assertEquals(Integer.valueOf(1), coordinator.getReplicationFactor(dataSegment.getId())));
+
     coordinator.stop();
     leaderUnannouncerLatch.await();