From cc53ca8d69bf598abe12c4a7afbf994a8170d2aa Mon Sep 17 00:00:00 2001
From: Colin Patrick Mccabe <cmccabe@cloudera.com>
Date: Tue, 31 May 2016 17:37:52 -0700
Subject: [PATCH] HDFS-9466.
 TestShortCircuitCache#testDataXceiverCleansUpSlotsOnFailure is flaky
 (Wei-Chiu Chuang via cmccabe)

(cherry picked from commit c7921c9bddb79c9db5059b6c3f7a3a586a3cd95b)
(cherry picked from commit 49c0d9b325a21fc0a05da58bb9d4c76d82181a5d)
---
 .../server/datanode/ShortCircuitRegistry.java |  6 ++--
 .../shortcircuit/TestShortCircuitCache.java   | 34 +++++++++----------
 2 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/hadoop-hdfs-project/hadoop-hdfs/src/main/java/org/apache/hadoop/hdfs/server/datanode/ShortCircuitRegistry.java b/hadoop-hdfs-project/hadoop-hdfs/src/main/java/org/apache/hadoop/hdfs/server/datanode/ShortCircuitRegistry.java
index 52856af2c46..68ef24e46d1 100644
--- a/hadoop-hdfs-project/hadoop-hdfs/src/main/java/org/apache/hadoop/hdfs/server/datanode/ShortCircuitRegistry.java
+++ b/hadoop-hdfs-project/hadoop-hdfs/src/main/java/org/apache/hadoop/hdfs/server/datanode/ShortCircuitRegistry.java
@@ -394,12 +394,12 @@ public void shutdown() {
   }
 
   public static interface Visitor {
-    void accept(HashMap<ShmId, RegisteredShm> segments,
+    boolean accept(HashMap<ShmId, RegisteredShm> segments,
                 HashMultimap<ExtendedBlockId, Slot> slots);
   }
 
   @VisibleForTesting
-  public synchronized void visit(Visitor visitor) {
-    visitor.accept(segments, slots);
+  public synchronized boolean visit(Visitor visitor) {
+    return visitor.accept(segments, slots);
   }
 }
diff --git a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/shortcircuit/TestShortCircuitCache.java b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/shortcircuit/TestShortCircuitCache.java
index f78861373cb..ac14438c633 100644
--- a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/shortcircuit/TestShortCircuitCache.java
+++ b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/shortcircuit/TestShortCircuitCache.java
@@ -32,6 +32,7 @@
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.concurrent.TimeoutException;
 
 import org.apache.commons.lang.mutable.MutableBoolean;
 import org.apache.commons.logging.Log;
@@ -624,15 +625,22 @@ public void visit(HashMap<DatanodeInfo,
   }
 
   static private void checkNumberOfSegmentsAndSlots(final int expectedSegments,
-        final int expectedSlots, ShortCircuitRegistry registry) {
-    registry.visit(new ShortCircuitRegistry.Visitor() {
+        final int expectedSlots, final ShortCircuitRegistry registry)
+  throws InterruptedException, TimeoutException {
+    GenericTestUtils.waitFor(new Supplier<Boolean>() {
       @Override
-      public void accept(HashMap<ShmId, RegisteredShm> segments,
-                         HashMultimap<ExtendedBlockId, Slot> slots) {
-        Assert.assertEquals(expectedSegments, segments.size());
-        Assert.assertEquals(expectedSlots, slots.size());
+      public Boolean get() {
+        return registry.visit(new ShortCircuitRegistry.Visitor() {
+          @Override
+          public boolean accept(HashMap<ShmId, RegisteredShm> segments,
+              HashMultimap<ExtendedBlockId, Slot> slots) {
+            return (expectedSegments == segments.size()) &&
+                (expectedSlots == slots.size());
+          }
+        });
       }
-    });
+    }, 100, 10000);
+
   }
 
   public static class TestCleanupFailureInjector
@@ -774,16 +782,8 @@ public void testPreReceiptVerificationDfsClientCanDoScr() throws Exception {
     DFSTestUtil.createFile(fs, TEST_PATH2, 4096, (short)1, 0xFADE2);
     DFSTestUtil.readFileBuffer(fs, TEST_PATH1);
     DFSTestUtil.readFileBuffer(fs, TEST_PATH2);
-    ShortCircuitRegistry registry =
-        cluster.getDataNodes().get(0).getShortCircuitRegistry();
-    registry.visit(new ShortCircuitRegistry.Visitor() {
-      @Override
-      public void accept(HashMap<ShmId, RegisteredShm> segments,
-                         HashMultimap<ExtendedBlockId, Slot> slots) {
-        Assert.assertEquals(1, segments.size());
-        Assert.assertEquals(2, slots.size());
-      }
-    });
+    checkNumberOfSegmentsAndSlots(1, 2,
+        cluster.getDataNodes().get(0).getShortCircuitRegistry());
     cluster.shutdown();
     sockDir.close();
   }