fix tests

Author: smalyshev

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/ExchangeService.java

@@ -185,6 +185,7 @@ public void finishSessionEarly(String sessionId, ActionListener<Void> listener)
  ExchangeSourceHandler exchangeSource = removeExchangeSourceHandler(sessionId);
  if (exchangeSource != null) {
  exchangeSource.finishEarly(false, listener);
+ exchangeSource.onFinishEarly();
  } else {
  listener.onResponse(null);
  }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/ExchangeSinkHandler.java

@@ -102,7 +102,7 @@ public void fetchPageAsync(boolean sourceFinished, ActionListener<ExchangeRespon
  buffer.finish(true);
  var subSource = source.get();
  if (subSource != null) {
- subSource.finishEarly(true, ActionListener.noop());
+ subSource.finishEarly(false, ActionListener.noop());
  }
  }
  listeners.add(listener);

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/ExchangeSourceHandler.java

@@ -317,9 +317,6 @@ public Releasable addEmptySink() {
  * @param drainingPages whether to discard pages already fetched in the exchange
  */
  public void finishEarly(boolean drainingPages, ActionListener<Void> listener) {
- if (finishEarlyHandler != null) {
- finishEarlyHandler.run();
- }
  buffer.finish(drainingPages);
  if (remoteSinks.isEmpty()) {
  listener.onResponse(null);
@@ -332,6 +329,15 @@ public void finishEarly(boolean drainingPages, ActionListener<Void> listener) {
  }
  }
 
+ /**
+ * Calls the handler signifying that the request has been completed prematurely.
+ */
+ public void onFinishEarly() {
+ if (finishEarlyHandler != null) {
+ finishEarlyHandler.run();
+ }
+ }
+
  private static class PendingInstances {
  private final AtomicInteger instances = new AtomicInteger();
  private final SubscribableListener<Void> completion = new SubscribableListener<>();

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponseTests.java

@@ -518,30 +518,65 @@ static EsqlQueryResponse fromXContent(XContentParser parser) {
  }
  }
 
+ public static int clusterDetailsSize(int numClusters) {
+ /* Example:
+ "_clusters" : {
+ "total" : 2,
+ "successful" : 2,
+ "running" : 0,
+ "skipped" : 0,
+ "partial" : 0,
+ "failed" : 0,
+ "is_partial" : false,
+ "details" : {
+ "(local)" : {
+ "status" : "successful",
+ "indices" : "logs-1",
+ "took" : 4444,
+ "_shards" : {
+ "total" : 10,
+ "successful" : 10,
+ "skipped" : 3,
+ "failed" : 0
+ }
+ },
+ "remote1" : {
+ "status" : "successful",
+ "indices" : "remote1:logs-1",
+ "took" : 4999,
+ "_shards" : {
+ "total" : 12,
+ "successful" : 12,
+ "skipped" : 5,
+ "failed" : 0
+ }
+ }
+ */
+ return numClusters * 4 + 7;
+ }
+
  public void testChunkResponseSizeColumnar() {
- int sizeClusterDetails = 14;
  try (EsqlQueryResponse resp = randomResponse(true, null)) {
  int columnCount = resp.pages().get(0).getBlockCount();
  int bodySize = resp.pages().stream().mapToInt(p -> p.getPositionCount() * p.getBlockCount()).sum() + columnCount * 2;
- assertChunkCount(resp, r -> 5 + sizeClusterDetails + bodySize);
+ assertChunkCount(resp, r -> 5 + clusterDetailsSize(resp.getExecutionInfo().clusterInfo.size()) + bodySize);
  }
 
  try (EsqlQueryResponse resp = randomResponseAsync(true, null, true)) {
  int columnCount = resp.pages().get(0).getBlockCount();
  int bodySize = resp.pages().stream().mapToInt(p -> p.getPositionCount() * p.getBlockCount()).sum() + columnCount * 2;
- assertChunkCount(resp, r -> 7 + sizeClusterDetails + bodySize); // is_running
+ assertChunkCount(resp, r -> 7 + clusterDetailsSize(resp.getExecutionInfo().clusterInfo.size()) + bodySize); // is_running
  }
  }
 
  public void testChunkResponseSizeRows() {
- int sizeClusterDetails = 14;
  try (EsqlQueryResponse resp = randomResponse(false, null)) {
  int bodySize = resp.pages().stream().mapToInt(Page::getPositionCount).sum();
- assertChunkCount(resp, r -> 5 + sizeClusterDetails + bodySize);
+ assertChunkCount(resp, r -> 5 + clusterDetailsSize(resp.getExecutionInfo().clusterInfo.size()) + bodySize);
  }
  try (EsqlQueryResponse resp = randomResponseAsync(false, null, true)) {
  int bodySize = resp.pages().stream().mapToInt(Page::getPositionCount).sum();
- assertChunkCount(resp, r -> 7 + sizeClusterDetails + bodySize);
+ assertChunkCount(resp, r -> 7 + clusterDetailsSize(resp.getExecutionInfo().clusterInfo.size()) + bodySize);
  }
  }
 

x-pack/plugin/security/qa/multi-cluster/src/javaRestTest/java/org/elasticsearch/xpack/remotecluster/CrossClusterEsqlRCS1MissingIndicesIT.java

@@ -79,6 +79,7 @@ record ExpectedCluster(String clusterAlias, String indexExpression, String statu
  void assertExpectedClustersForMissingIndicesTests(Map<String, Object> responseMap, List<ExpectedCluster> expected) {
  Map<String, ?> clusters = (Map<String, ?>) responseMap.get("_clusters");
  assertThat((int) responseMap.get("took"), greaterThan(0));
+ assertThat((boolean) clusters.get("is_partial"), is(false));
 
  Map<String, ?> detailsMap = (Map<String, ?>) clusters.get("details");
  assertThat(detailsMap.size(), is(expected.size()));

x-pack/plugin/security/qa/multi-cluster/src/javaRestTest/java/org/elasticsearch/xpack/remotecluster/CrossClusterEsqlRCS1UnavailableRemotesIT.java

@@ -101,6 +101,7 @@ private void clusterShutDownWithRandomSkipUnavailable() throws Exception {
  assertThat((int) clusters.get("skipped"), is(0));
  assertThat((int) clusters.get("partial"), is(0));
  assertThat((int) clusters.get("failed"), is(0));
+ assertThat((boolean) clusters.get("is_partial"), is(false));
 
  assertThat(clusterDetails.size(), is(2));
  assertThat((int) localClusterDetails.get("took"), greaterThan(0));

x-pack/plugin/security/qa/multi-cluster/src/javaRestTest/java/org/elasticsearch/xpack/remotecluster/CrossClusterEsqlRCS2UnavailableRemotesIT.java

@@ -119,6 +119,7 @@ private void clusterShutDownWithRandomSkipUnavailable() throws Exception {
  assertThat((int) clusters.get("skipped"), is(0));
  assertThat((int) clusters.get("partial"), is(0));
  assertThat((int) clusters.get("failed"), is(0));
+ assertThat((boolean) clusters.get("is_partial"), is(false));
 
  assertThat(clusterDetails.size(), is(2));
  assertThat((int) localClusterDetails.get("took"), greaterThan(0));