From fc087013aed3f5bb86beb4c9edf31eca9440bf9a Mon Sep 17 00:00:00 2001
From: Jihoon Son <jihoonson@apache.org>
Date: Thu, 6 Jul 2017 00:57:38 +0900
Subject: [PATCH] Fix a bug in DruidCluster.getAllServers() (#4500)

* Fix a bug in getAllServers

* Change to old style
---
 .../io/druid/server/coordinator/DruidCluster.java    | 12 +++++++++---
 .../druid/server/coordinator/DruidClusterTest.java   |  7 +++++++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/server/src/main/java/io/druid/server/coordinator/DruidCluster.java b/server/src/main/java/io/druid/server/coordinator/DruidCluster.java
index 20561bee8550..014c42ba91c4 100644
--- a/server/src/main/java/io/druid/server/coordinator/DruidCluster.java
+++ b/server/src/main/java/io/druid/server/coordinator/DruidCluster.java
@@ -26,9 +26,11 @@
 import io.druid.java.util.common.IAE;
 
 import javax.annotation.Nullable;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -113,9 +115,13 @@ public MinMaxPriorityQueue<ServerHolder> getHistoricalsByTier(String tier)
 
   public Collection<ServerHolder> getAllServers()
   {
-    return historicals.values().stream()
-                      .flatMap(Collection::stream)
-                      .collect(() -> realtimes, Set::add, Set::addAll);
+    final int historicalSize = historicals.values().stream().mapToInt(Collection::size).sum();
+    final int realtimeSize = realtimes.size();
+    final List<ServerHolder> allServers = new ArrayList<>(historicalSize + realtimeSize);
+
+    historicals.values().forEach(allServers::addAll);
+    allServers.addAll(realtimes);
+    return allServers;
   }
 
   public Iterable<MinMaxPriorityQueue<ServerHolder>> getSortedHistoricalsByTier()
diff --git a/server/src/test/java/io/druid/server/coordinator/DruidClusterTest.java b/server/src/test/java/io/druid/server/coordinator/DruidClusterTest.java
index d91675f9e42f..9d001146c2dc 100644
--- a/server/src/test/java/io/druid/server/coordinator/DruidClusterTest.java
+++ b/server/src/test/java/io/druid/server/coordinator/DruidClusterTest.java
@@ -38,6 +38,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 public class DruidClusterTest
@@ -184,6 +185,9 @@ public void testGetAllServers()
   {
     cluster.add(newRealtime);
     cluster.add(newHistorical);
+    final Set<ServerHolder> expectedRealtimes = cluster.getRealtimes();
+    final Map<String, MinMaxPriorityQueue<ServerHolder>> expectedHistoricals = cluster.getHistoricals();
+
     final Collection<ServerHolder> allServers = cluster.getAllServers();
     Assert.assertEquals(4, allServers.size());
     Assert.assertTrue(allServers.containsAll(cluster.getRealtimes()));
@@ -192,6 +196,9 @@ public void testGetAllServers()
             cluster.getHistoricals().values().stream().flatMap(Collection::stream).collect(Collectors.toList())
         )
     );
+
+    Assert.assertEquals(expectedHistoricals, cluster.getHistoricals());
+    Assert.assertEquals(expectedRealtimes, cluster.getRealtimes());
   }
 
   @Test