From 54a8fb827df6ca2d7b0f1e86656d0cd5a978f273 Mon Sep 17 00:00:00 2001
From: Maytas Monsereenusorn <maytasm@apache.org>
Date: Wed, 8 Jul 2020 20:03:52 -0700
Subject: [PATCH] Fix flaky tests in DruidCoordinatorTest (#10157)

* Fix flaky tests in DruidCoordinatorTest

* Imporve fail msg

* Fix flaky tests in DruidCoordinatorTest
---
 .../coordinator/DruidCoordinatorTest.java      | 18 ++++++------------
 1 file changed, 6 insertions(+), 12 deletions(-)

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
index 0077fc1123a..51308cbd8fa 100644
--- a/server/src/test/java/org/apache/druid/server/coordinator/DruidCoordinatorTest.java
+++ b/server/src/test/java/org/apache/druid/server/coordinator/DruidCoordinatorTest.java
@@ -103,12 +103,10 @@ public class DruidCoordinatorTest extends CuratorTestBase
   private ObjectMapper objectMapper;
   private DruidNode druidNode;
   private LatchableServiceEmitter serviceEmitter = new LatchableServiceEmitter();
-  private boolean serverAddedCountExpected = true;
 
   @Before
   public void setUp() throws Exception
   {
-    serverAddedCountExpected = true;
     druidServer = EasyMock.createMock(DruidServer.class);
     serverInventoryView = EasyMock.createMock(SingleServerInventoryView.class);
     segmentsMetadataManager = EasyMock.createNiceMock(SegmentsMetadataManager.class);
@@ -385,7 +383,6 @@ public class DruidCoordinatorTest extends CuratorTestBase
         druidServer
     );
     assignSegmentLatch.await();
-    Assert.assertTrue(serverAddedCountExpected);
 
     final CountDownLatch coordinatorRunLatch = new CountDownLatch(2);
     serviceEmitter.latch = coordinatorRunLatch;
@@ -486,7 +483,6 @@ public class DruidCoordinatorTest extends CuratorTestBase
     final CountDownLatch assignSegmentLatchCold = createCountDownLatchAndSetPathChildrenCacheListenerWithLatch(1, pathChildrenCacheCold, dataSegments, coldServer);
     assignSegmentLatchHot.await();
     assignSegmentLatchCold.await();
-    Assert.assertTrue(serverAddedCountExpected);
 
     final CountDownLatch coordinatorRunLatch = new CountDownLatch(2);
     serviceEmitter.latch = coordinatorRunLatch;
@@ -646,7 +642,6 @@ public class DruidCoordinatorTest extends CuratorTestBase
     assignSegmentLatchBroker1.await();
     assignSegmentLatchBroker2.await();
     assignSegmentLatchPeon.await();
-    Assert.assertTrue(serverAddedCountExpected);
 
     final CountDownLatch coordinatorRunLatch = new CountDownLatch(2);
     serviceEmitter.latch = coordinatorRunLatch;
@@ -679,16 +674,15 @@ public class DruidCoordinatorTest extends CuratorTestBase
     pathChildrenCache.getListenable().addListener(
         (CuratorFramework client, PathChildrenCacheEvent event) -> {
           if (CuratorUtils.isChildAdded(event)) {
-            if (countDownLatch.getCount() > 0) {
-              DataSegment segment = findSegmentRelatedToCuratorEvent(segments, event);
-              if (segment != null) {
+            DataSegment segment = findSegmentRelatedToCuratorEvent(segments, event);
+            if (segment != null && server.getSegment(segment.getId()) == null) {
+              if (countDownLatch.getCount() > 0) {
                 server.addDataSegment(segment);
                 curator.delete().guaranteed().forPath(event.getData().getPath());
+                countDownLatch.countDown();
+              } else {
+                Assert.fail("The segment path " + event.getData().getPath() + " is not expected");
               }
-              countDownLatch.countDown();
-            } else {
-              // The segment is assigned to the server more times than expected
-              serverAddedCountExpected = false;
             }
           }
         }