PIT relocation POC

[CI] Auto commit changes from spotless

Author: cbuescher

server/src/main/java/org/elasticsearch/action/ResolvedIndices.java

@@ -61,6 +61,10 @@ public class ResolvedIndices {
         this(remoteClusterIndices, localIndices, localIndexMetadata, null);
     }
 
+    public String toString() {
+        return "ResolvedIndices{" + "local=" + localIndices + ", remote=" + remoteClusterIndices + '}';
+    }
+
     /**
      * Get the remote cluster indices, structured as a map where the key is the remote cluster alias.
      * <br/>

server/src/main/java/org/elasticsearch/action/search/ClearScrollController.java

@@ -8,20 +8,27 @@
  */
 package org.elasticsearch.action.search;
 
+import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.support.RefCountingRunnable;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.node.DiscoveryNodes;
+import org.elasticsearch.cluster.project.ProjectResolver;
+import org.elasticsearch.cluster.routing.IndexShardRoutingTable;
+import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.util.concurrent.ListenableFuture;
+import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.transport.Transport;
 import org.elasticsearch.transport.TransportResponse;
 
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -38,6 +45,7 @@ public final class ClearScrollController implements Runnable {
     private final AtomicBoolean hasFailed = new AtomicBoolean(false);
     private final AtomicInteger freedSearchContexts = new AtomicInteger(0);
     private final Logger logger;
+    private static final Logger staticLogger = LogManager.getLogger(ClearScrollController.class);
     private final Runnable runner;
 
     ClearScrollController(
@@ -148,12 +156,15 @@ private void finish() {
      * Closes the given context id and reports the number of freed contexts via the listener
      */
     public static void closeContexts(
-        DiscoveryNodes nodes,
+        ClusterService clusterService,
+        ProjectResolver projectResolver,
         SearchTransportService searchTransportService,
-        Collection<SearchContextIdForNode> contextIds,
+        Map<ShardId, SearchContextIdForNode> shards,
         ActionListener<Integer> listener
     ) {
-        final Set<String> clusters = contextIds.stream()
+        DiscoveryNodes nodes = clusterService.state().nodes();
+        final Set<String> clusters = shards.values()
+            .stream()
             .map(SearchContextIdForNode::getClusterAlias)
             .filter(clusterAlias -> Strings.isEmpty(clusterAlias) == false)
             .collect(Collectors.toSet());
@@ -166,16 +177,34 @@ public static void closeContexts(
         lookupListener.addListener(listener.delegateFailure((l, nodeLookup) -> {
             final var successes = new AtomicInteger();
             try (RefCountingRunnable refs = new RefCountingRunnable(() -> l.onResponse(successes.get()))) {
-                for (SearchContextIdForNode contextId : contextIds) {
+                for (Entry<ShardId, SearchContextIdForNode> entry : shards.entrySet()) {
+                    var contextId = entry.getValue();
                     if (contextId.getNode() == null) {
                         // the shard was missing when creating the PIT, ignore.
                         continue;
                     }
                     final DiscoveryNode node = nodeLookup.apply(contextId.getClusterAlias(), contextId.getNode());
+
+                    Set<DiscoveryNode> targetNodes;
                     if (node != null) {
+                        targetNodes = Collections.singleton(node);
+                    } else {
+                        staticLogger.info("---> missing node when closing context: " + contextId.getNode());
+                        // TODO we won't be able to use this with remote clusters
+                        IndexShardRoutingTable indexShardRoutingTable = clusterService.state()
+                            .routingTable(projectResolver.getProjectId())
+                            .shardRoutingTable(entry.getKey());
+                        targetNodes = indexShardRoutingTable.assignedUnpromotableShards()
+                            .stream()
+                            .map(ShardRouting::currentNodeId)
+                            .map(nodeId -> nodeLookup.apply(contextId.getClusterAlias(), nodeId))
+                            .collect(Collectors.toSet());
+                        staticLogger.info("---> trying alternative nodes to close context: " + targetNodes);
+                    }
+                    for (DiscoveryNode targetNode : targetNodes) {
                         try {
                             searchTransportService.sendFreeContext(
-                                searchTransportService.getConnection(contextId.getClusterAlias(), node),
+                                searchTransportService.getConnection(contextId.getClusterAlias(), targetNode),
                                 contextId.getSearchContextId(),
                                 refs.acquireListener().map(r -> {
                                     if (r.isFreed()) {

server/src/main/java/org/elasticsearch/action/search/PITHelper.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.action.search;
+
+import org.elasticsearch.TransportVersion;
+import org.elasticsearch.common.bytes.BytesArray;
+import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.common.util.Maps;
+import org.elasticsearch.index.shard.ShardId;
+
+import java.io.IOException;
+import java.util.Base64;
+import java.util.Collections;
+import java.util.Map;
+
+public class PITHelper {
+
+    public static SearchContextId decodePITId(String id) throws IOException {
+        return decodePITId(new BytesArray(Base64.getUrlDecoder().decode(id)));
+    }
+
+    public static SearchContextId decodePITId(BytesReference id) throws IOException {
+        try (var in = id.streamInput()) {
+            final TransportVersion version = TransportVersion.readVersion(in);
+            in.setTransportVersion(version);
+            final Map<ShardId, SearchContextIdForNode> shards = Collections.unmodifiableMap(
+                in.readCollection(Maps::newHashMapWithExpectedSize, (i, map) -> map.put(new ShardId(in), new SearchContextIdForNode(in)))
+            );
+            return new SearchContextId(shards, Collections.emptyMap());
+        } catch (IOException e) {
+            assert false : e;
+            throw new IllegalArgumentException(e);
+        }
+    }
+}

server/src/main/java/org/elasticsearch/action/search/TransportClosePointInTimeAction.java

@@ -13,45 +13,51 @@
 import org.elasticsearch.action.ActionType;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.HandledTransportAction;
+import org.elasticsearch.cluster.project.ProjectResolver;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
+import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.injection.guice.Inject;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.transport.TransportService;
 
-import java.util.Collection;
+import java.util.Map;
 
 public class TransportClosePointInTimeAction extends HandledTransportAction<ClosePointInTimeRequest, ClosePointInTimeResponse> {
 
     public static final ActionType<ClosePointInTimeResponse> TYPE = new ActionType<>("indices:data/read/close_point_in_time");
     private final ClusterService clusterService;
     private final SearchTransportService searchTransportService;
     private final NamedWriteableRegistry namedWriteableRegistry;
+    private final ProjectResolver projectResolver;
 
     @Inject
     public TransportClosePointInTimeAction(
         TransportService transportService,
         ClusterService clusterService,
         ActionFilters actionFilters,
         SearchTransportService searchTransportService,
-        NamedWriteableRegistry namedWriteableRegistry
+        NamedWriteableRegistry namedWriteableRegistry,
+        ProjectResolver projectResolver
     ) {
         super(TYPE.name(), transportService, actionFilters, ClosePointInTimeRequest::new, EsExecutors.DIRECT_EXECUTOR_SERVICE);
         this.clusterService = clusterService;
         this.searchTransportService = searchTransportService;
         this.namedWriteableRegistry = namedWriteableRegistry;
+        this.projectResolver = projectResolver;
     }
 
     @Override
     protected void doExecute(Task task, ClosePointInTimeRequest request, ActionListener<ClosePointInTimeResponse> listener) {
         final SearchContextId searchContextId = SearchContextId.decode(namedWriteableRegistry, request.getId());
-        final Collection<SearchContextIdForNode> contextIds = searchContextId.shards().values();
+        Map<ShardId, SearchContextIdForNode> shards = searchContextId.shards();
         ClearScrollController.closeContexts(
-            clusterService.state().nodes(),
+            clusterService,
+            projectResolver,
             searchTransportService,
-            contextIds,
-            listener.map(freed -> new ClosePointInTimeResponse(freed == contextIds.size(), freed))
+            shards,
+            listener.map(freed -> new ClosePointInTimeResponse(freed == shards.size(), freed))
         );
     }
 }

server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java

@@ -371,6 +371,7 @@ void executeRequest(
             );
             frozenIndexCheck(resolvedIndices);
         }
+        logger.info("Executing search request on node [{}] with indices [{}]", clusterService.getNodeName(), resolvedIndices);
 
         ActionListener<SearchRequest> rewriteListener = listener.delegateFailureAndWrap((delegate, rewritten) -> {
             if (ccsCheckCompatibility) {
@@ -1313,6 +1314,7 @@ private void executeSearch(
         SearchResponse.Clusters clusters,
         SearchPhaseProvider searchPhaseProvider
     ) {
+        logger.info("Executing search locally.");
         if (searchRequest.allowPartialSearchResults() == null) {
             // No user preference defined in search request - apply cluster service default
             searchRequest.allowPartialSearchResults(searchService.defaultAllowPartialSearchResults());
@@ -1905,11 +1907,9 @@ static List<SearchShardIterator> getLocalShardsIteratorFromPointInTime(
                         if (projectState.cluster().nodes().nodeExists(perNode.getNode())) {
                             targetNodes.add(perNode.getNode());
                         }
-                        if (perNode.getSearchContextId().getSearcherId() != null) {
-                            for (ShardRouting shard : shards) {
-                                if (shard.currentNodeId().equals(perNode.getNode()) == false) {
-                                    targetNodes.add(shard.currentNodeId());
-                                }
+                        for (ShardRouting shard : shards) {
+                            if (shard.currentNodeId().equals(perNode.getNode()) == false) {
+                                targetNodes.add(shard.currentNodeId());
                             }
                         }
                     } catch (IndexNotFoundException | ShardNotFoundException e) {

server/src/main/java/org/elasticsearch/index/engine/ReadOnlyEngine.java

@@ -152,7 +152,7 @@ public ReadOnlyEngine(
      * as the search id because the commit id changes whenever IndexWriter#commit is called although the segment files stay unchanged.
      * Any recovery except the local recovery performs IndexWriter#commit to generate a new translog uuid or history_uuid.
      */
-    static String generateSearcherId(SegmentInfos sis) {
+    public static String generateSearcherId(SegmentInfos sis) {
         final MessageDigest md = MessageDigests.sha256();
         for (SegmentCommitInfo si : sis) {
             final byte[] segmentId = si.getId();

server/src/main/java/org/elasticsearch/search/SearchContextMissingException.java

@@ -22,7 +22,7 @@ public class SearchContextMissingException extends ElasticsearchException {
     private final ShardSearchContextId contextId;
 
     public SearchContextMissingException(ShardSearchContextId contextId) {
-        super("No search context found for id [" + contextId.getId() + "]");
+        super("No search context found for id [" + contextId.getSessionId() + "/" + contextId.getId() + "]");
         this.contextId = contextId;
     }