fix TestShuffleChannelHandler

Author: stoty

hadoop-mapreduce-project/hadoop-mapreduce-client/hadoop-mapreduce-client-shuffle/src/test/java/org/apache/hadoop/mapred/TestShuffleChannelHandler.java

@@ -44,6 +44,7 @@
 import io.netty.handler.ssl.SslContextBuilder;
 import io.netty.handler.ssl.SslHandler;
 import io.netty.handler.stream.ChunkedWriteHandler;
+import io.netty.util.ReferenceCounted;
 import io.netty.util.concurrent.GlobalEventExecutor;
 
 import java.io.ByteArrayOutputStream;
@@ -56,11 +57,13 @@
 import java.security.cert.X509Certificate;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 
 import javax.crypto.SecretKey;
@@ -115,6 +118,7 @@ public void testGetMapsChunkedFileSSl() throws Exception {
     final LinkedList<Object> unencryptedMessages = new LinkedList<>();
     final EmbeddedChannel shuffle = t.createShuffleHandlerSSL(unencryptedMessages);
     t.testGetAllAttemptsForReduce0NoKeepAlive(unencryptedMessages, shuffle);
+    drainChannel(shuffle);
   }
 
   @Test
@@ -194,6 +198,8 @@ public void testIncompatibleShuffleVersion() {
         actual.toString());
 
     assertFalse(shuffle.isActive(), "closed"); // known-issue
+    tryRelease(actual);
+    drainChannel(decoder);
   }
 
   @Test
@@ -217,6 +223,9 @@ public void testInvalidMapNoIndexFile() {
         actual.toString());
 
     assertFalse(shuffle.isActive(), "closed");
+
+    tryRelease(actual);
+    drainChannel(decoder);
   }
 
   @Test
@@ -243,9 +252,31 @@ public void testInvalidMapNoDataFile() {
     assertEquals(getExpectedHttpResponse(HttpResponseStatus.INTERNAL_SERVER_ERROR).toString(),
         actual.toString());
 
+    tryRelease(actual);
+    drainChannel(decoder);
+    
     assertFalse(shuffle.isActive(), "closed");
   }
 
+  private void drainChannel(EmbeddedChannel ch) {
+    Object o;
+    while((o = ch.readInbound())!=null) {
+      tryRelease(o);
+    }
+    while((o = ch.readOutbound())!=null) {
+      tryRelease(o);
+    }
+  }
+
+  private void tryRelease(Object obj) {
+    if (obj instanceof ReferenceCounted) {
+      ReferenceCounted bb = (ReferenceCounted) obj;
+      if (bb.refCnt() > 0) {
+        bb.release(bb.refCnt());
+      }
+    }
+  }
+
   private DefaultHttpResponse getExpectedHttpResponse(HttpResponseStatus status) {
     DefaultHttpResponse response = new DefaultHttpResponse(HTTP_1_1, status);
     response.headers().set(CONTENT_TYPE, "text/plain; charset=UTF-8");
@@ -366,7 +397,7 @@ private void testGetAllAttemptsForReduce0NoKeepAlive(
     }
 
     private void testKeepAlive(java.util.Queue<Object> messages,
-                               EmbeddedChannel shuffle) throws IOException {
+                               EmbeddedChannel shuffle) throws IOException, InterruptedException, ExecutionException {
       final FullHttpRequest req1 = createRequest(
           getUri(TEST_JOB_ID, 0, Collections.singletonList(TEST_ATTEMPT_1), true));
       shuffle.writeInbound(req1);
@@ -375,6 +406,7 @@ private void testKeepAlive(java.util.Queue<Object> messages,
           getAttemptData(new Attempt(TEST_ATTEMPT_1, TEST_DATA_A))
       );
       assertTrue(shuffle.isActive(), "keep-alive");
+      drainChannel(shuffle);
       messages.clear();
 
       final FullHttpRequest req2 = createRequest(
@@ -385,6 +417,7 @@ private void testKeepAlive(java.util.Queue<Object> messages,
           getAttemptData(new Attempt(TEST_ATTEMPT_2, TEST_DATA_B))
       );
       assertTrue(shuffle.isActive(), "keep-alive");
+      drainChannel(shuffle);
       messages.clear();
 
       final FullHttpRequest req3 = createRequest(
@@ -395,6 +428,7 @@ private void testKeepAlive(java.util.Queue<Object> messages,
           getAttemptData(new Attempt(TEST_ATTEMPT_3, TEST_DATA_C))
       );
       assertFalse(shuffle.isActive(), "no keep-alive");
+      drainChannel(shuffle);
     }
 
     private ArrayList<ByteBuf> getAllAttemptsForReduce0() throws IOException {
@@ -442,6 +476,8 @@ private void assertResponse(java.util.Queue<Object> outboundMessages,
         }
 
         i++;
+        tryRelease(obj);
+        drainChannel(decodeChannel);
       }
 
       // This check is done after to have better debug logs on failure.

hadoop-mapreduce-project/hadoop-mapreduce-client/hadoop-mapreduce-client-shuffle/src/test/java/org/apache/hadoop/mapred/TestShuffleHandlerBase.java

@@ -84,6 +84,13 @@ public void setup() throws IOException {
 
   @AfterEach
   public void teardown() {
+    //Trigger GC so that we get the leak warnings early
+    System.gc();
+    try {
+      // Wait for logger to flush
+      Thread.sleep(1000);
+    } catch (InterruptedException e) {
+    }
     System.setOut(standardOut);
     System.out.print(outputStreamCaptor);
     // For this to work ch.qos.logback.classic is needed for some reason