From 91b598a3ae0a78aff90fa2a7178ce2d886d58470 Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <nhat.nguyen@elastic.co>
Date: Wed, 13 Jun 2018 20:49:50 -0400
Subject: [PATCH 1/2] TEST: getCapturedRequestsAndClear should be atomic

We might lose messages between getCapturedRequestsAndClear calls. This
commit makes sure that both getCapturedRequestsAndClear and
getCapturedRequestsByTargetNodeAndClear are atomic.
---
 .../test/transport/CapturingTransport.java    | 38 +++++++++++--------
 1 file changed, 22 insertions(+), 16 deletions(-)

diff --git a/test/framework/src/main/java/org/elasticsearch/test/transport/CapturingTransport.java b/test/framework/src/main/java/org/elasticsearch/test/transport/CapturingTransport.java
index 81fc934ca6d7e..a7a0eb446e8fb 100644
--- a/test/framework/src/main/java/org/elasticsearch/test/transport/CapturingTransport.java
+++ b/test/framework/src/main/java/org/elasticsearch/test/transport/CapturingTransport.java
@@ -46,6 +46,7 @@
 import java.io.UncheckedIOException;
 import java.net.UnknownHostException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -94,9 +95,23 @@ public CapturedRequest[] capturedRequests() {
      * @return the captured requests
      */
     public CapturedRequest[] getCapturedRequestsAndClear() {
-        CapturedRequest[] capturedRequests = capturedRequests();
-        clear();
-        return capturedRequests;
+        List<CapturedRequest> requests = new ArrayList<>(capturedRequests.size());
+        capturedRequests.drainTo(requests);
+        return requests.toArray(new CapturedRequest[0]);
+    }
+
+    private Map<String, List<CapturedRequest>> groupRequestsByTargetNode(Collection<CapturedRequest> requests) {
+        Map<String, List<CapturedRequest>> result = new HashMap<>();
+        for (CapturedRequest request : requests) {
+            result.compute(request.node.getId(), (k, group) -> {
+                if (group == null) {
+                    group = new ArrayList<>();
+                }
+                group.add(request);
+                return group;
+            });
+        }
+        return result;
     }
 
     /**
@@ -104,16 +119,7 @@ public CapturedRequest[] getCapturedRequestsAndClear() {
      * Doesn't clear the captured request list. See {@link #clear()}
      */
     public Map<String, List<CapturedRequest>> capturedRequestsByTargetNode() {
-        Map<String, List<CapturedRequest>> map = new HashMap<>();
-        for (CapturedRequest request : capturedRequests) {
-            List<CapturedRequest> nodeList = map.get(request.node.getId());
-            if (nodeList == null) {
-                nodeList = new ArrayList<>();
-                map.put(request.node.getId(), nodeList);
-            }
-            nodeList.add(request);
-        }
-        return map;
+        return groupRequestsByTargetNode(capturedRequests);
     }
 
     /**
@@ -125,9 +131,9 @@ public Map<String, List<CapturedRequest>> capturedRequestsByTargetNode() {
      * @return the captured requests grouped by target node
      */
     public Map<String, List<CapturedRequest>> getCapturedRequestsByTargetNodeAndClear() {
-        Map<String, List<CapturedRequest>> map = capturedRequestsByTargetNode();
-        clear();
-        return map;
+        List<CapturedRequest> requests = new ArrayList<>(capturedRequests.size());
+        capturedRequests.drainTo(requests);
+        return groupRequestsByTargetNode(requests);
     }
 
     /** clears captured requests */

From cda7bb912fdc2aa34876f109212097c566849b26 Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <nhat.nguyen@elastic.co>
Date: Thu, 14 Jun 2018 17:44:51 -0400
Subject: [PATCH 2/2] compute -> computeIfAbsent

---
 .../elasticsearch/test/transport/CapturingTransport.java  | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/test/framework/src/main/java/org/elasticsearch/test/transport/CapturingTransport.java b/test/framework/src/main/java/org/elasticsearch/test/transport/CapturingTransport.java
index a7a0eb446e8fb..318c70c2933d8 100644
--- a/test/framework/src/main/java/org/elasticsearch/test/transport/CapturingTransport.java
+++ b/test/framework/src/main/java/org/elasticsearch/test/transport/CapturingTransport.java
@@ -103,13 +103,7 @@ public CapturedRequest[] getCapturedRequestsAndClear() {
     private Map<String, List<CapturedRequest>> groupRequestsByTargetNode(Collection<CapturedRequest> requests) {
         Map<String, List<CapturedRequest>> result = new HashMap<>();
         for (CapturedRequest request : requests) {
-            result.compute(request.node.getId(), (k, group) -> {
-                if (group == null) {
-                    group = new ArrayList<>();
-                }
-                group.add(request);
-                return group;
-            });
+            result.computeIfAbsent(request.node.getId(), node -> new ArrayList<>()).add(request);
         }
         return result;
     }