Fix tests

elastic · Oct 4, 2023 · 63e6473 · 63e6473
1 parent c91e1f2
commit 63e6473
Showing 1 changed file with 11 additions and 4 deletions.
diff --git a/...mpute/src/test/java/org/elasticsearch/compute/operator/exchange/ExchangeServiceTests.java b/...mpute/src/test/java/org/elasticsearch/compute/operator/exchange/ExchangeServiceTests.java
@@ -398,18 +398,25 @@ public void messageReceived(
             ) throws Exception {
                 FilterTransportChannel filterChannel = new FilterTransportChannel(channel) {
                     @Override
-                    public void sendResponse(TransportResponse response) throws IOException {
-                        ExchangeResponse exchangeResponse = (ExchangeResponse) response;
-                        Page page = exchangeResponse.takePage();
+                    public void sendResponse(TransportResponse transportResponse) throws IOException {
+                        ExchangeResponse origResp = (ExchangeResponse) transportResponse;
+                        Page page = origResp.takePage();
                         if (page != null) {
                             IntBlock block = page.getBlock(0);
                             for (int i = 0; i < block.getPositionCount(); i++) {
                                 if (block.getInt(i) == disconnectOnSeqNo) {
+                                    page.releaseBlocks();
                                     throw new IOException("page is too large");
                                 }
                             }
                         }
-                        super.sendResponse(response);
+                        ExchangeResponse newResp = new ExchangeResponse(page, origResp.finished());
+                        origResp.decRef();
+                        while (origResp.hasReferences()) {
+                            newResp.incRef();
+                            origResp.decRef();
+                        }
+                        super.sendResponse(newResp);
                     }
                 };
                 handler.messageReceived(request, filterChannel, task);