Skip to content

Commit e5b7137

Browse files
authored
TEST: getCapturedRequestsAndClear should be atomic (#31312)
We might lose messages between getCapturedRequestsAndClear calls. This commit makes sure that both getCapturedRequestsAndClear and getCapturedRequestsByTargetNodeAndClear are atomic.
1 parent 0324103 commit e5b7137

File tree

1 file changed

+16
-16
lines changed

1 file changed

+16
-16
lines changed

test/framework/src/main/java/org/elasticsearch/test/transport/CapturingTransport.java

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import java.io.UncheckedIOException;
4747
import java.net.UnknownHostException;
4848
import java.util.ArrayList;
49+
import java.util.Collection;
4950
import java.util.Collections;
5051
import java.util.HashMap;
5152
import java.util.List;
@@ -94,26 +95,25 @@ public CapturedRequest[] capturedRequests() {
9495
* @return the captured requests
9596
*/
9697
public CapturedRequest[] getCapturedRequestsAndClear() {
97-
CapturedRequest[] capturedRequests = capturedRequests();
98-
clear();
99-
return capturedRequests;
98+
List<CapturedRequest> requests = new ArrayList<>(capturedRequests.size());
99+
capturedRequests.drainTo(requests);
100+
return requests.toArray(new CapturedRequest[0]);
101+
}
102+
103+
private Map<String, List<CapturedRequest>> groupRequestsByTargetNode(Collection<CapturedRequest> requests) {
104+
Map<String, List<CapturedRequest>> result = new HashMap<>();
105+
for (CapturedRequest request : requests) {
106+
result.computeIfAbsent(request.node.getId(), node -> new ArrayList<>()).add(request);
107+
}
108+
return result;
100109
}
101110

102111
/**
103112
* returns all requests captured so far, grouped by target node.
104113
* Doesn't clear the captured request list. See {@link #clear()}
105114
*/
106115
public Map<String, List<CapturedRequest>> capturedRequestsByTargetNode() {
107-
Map<String, List<CapturedRequest>> map = new HashMap<>();
108-
for (CapturedRequest request : capturedRequests) {
109-
List<CapturedRequest> nodeList = map.get(request.node.getId());
110-
if (nodeList == null) {
111-
nodeList = new ArrayList<>();
112-
map.put(request.node.getId(), nodeList);
113-
}
114-
nodeList.add(request);
115-
}
116-
return map;
116+
return groupRequestsByTargetNode(capturedRequests);
117117
}
118118

119119
/**
@@ -125,9 +125,9 @@ public Map<String, List<CapturedRequest>> capturedRequestsByTargetNode() {
125125
* @return the captured requests grouped by target node
126126
*/
127127
public Map<String, List<CapturedRequest>> getCapturedRequestsByTargetNodeAndClear() {
128-
Map<String, List<CapturedRequest>> map = capturedRequestsByTargetNode();
129-
clear();
130-
return map;
128+
List<CapturedRequest> requests = new ArrayList<>(capturedRequests.size());
129+
capturedRequests.drainTo(requests);
130+
return groupRequestsByTargetNode(requests);
131131
}
132132

133133
/** clears captured requests */

0 commit comments

Comments
 (0)