Prevent closing wrong search contexts in compute service

Author: dnhatn

muted-tests.yml

@@ -393,9 +393,6 @@ tests:
 - class: org.elasticsearch.packaging.test.DockerTests
   method: test026InstallBundledRepositoryPluginsViaConfigFile
   issue: https://github.com/elastic/elasticsearch/issues/127158
-- class: org.elasticsearch.xpack.esql.plugin.DataNodeRequestSenderIT
-  method: testSearchWhileRelocating
-  issue: https://github.com/elastic/elasticsearch/issues/127188
 - class: org.elasticsearch.xpack.remotecluster.CrossClusterEsqlRCS2EnrichUnavailableRemotesIT
   method: testEsqlEnrichWithSkipUnavailable
   issue: https://github.com/elastic/elasticsearch/issues/127368

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeComputeHandler.java

@@ -328,8 +328,8 @@ private void acquireSearchContexts(
             }
             final var doAcquire = ActionRunnable.supply(listener, () -> {
                 final List<SearchContext> searchContexts = new ArrayList<>(targetShards.size());
-                SearchContext context = null;
                 for (IndexShard shard : targetShards) {
+                    SearchContext context = null;
                     try {
                         var aliasFilter = aliasFilters.getOrDefault(shard.shardId().getIndex(), AliasFilter.EMPTY);
                         var shardRequest = new ShardSearchRequest(
@@ -342,15 +342,16 @@ private void acquireSearchContexts(
                         // we need to limit the number of active search contexts here or in SearchService
                         context = searchService.createSearchContext(shardRequest, SearchService.NO_TIMEOUT);
                         context.preProcess();
-                        searchContexts.add(context);
                     } catch (Exception e) {
                         if (addShardLevelFailure(shard.shardId(), e)) {
                             IOUtils.close(context);
+                            continue;
                         } else {
                             IOUtils.closeWhileHandlingException(context, () -> IOUtils.close(searchContexts));
                             throw e;
                         }
                     }
+                    searchContexts.add(context);
                 }
                 return searchContexts;
             });