Adding re-writing of PIT id

Author: cbuescher

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

@@ -13,6 +13,7 @@
 import org.apache.lucene.util.SetOnce;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.TransportVersion;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.NoShardAvailableActionException;
 import org.elasticsearch.action.OriginalIndices;
@@ -93,6 +94,7 @@ abstract class AbstractSearchAsyncAction<Result extends SearchPhaseResult> exten
     private final Map<String, PendingExecutions> pendingExecutionsPerNode;
     private final AtomicBoolean requestCancelled = new AtomicBoolean();
     private final int skippedCount;
+    private final TransportVersion clusterMinTransportVersion;
 
     // protected for tests
     protected final SubscribableListener<Void> doneFuture = new SubscribableListener<>();
@@ -149,6 +151,7 @@ abstract class AbstractSearchAsyncAction<Result extends SearchPhaseResult> exten
         this.nodeIdToConnection = nodeIdToConnection;
         this.concreteIndexBoosts = concreteIndexBoosts;
         this.clusterStateVersion = clusterState.version();
+        this.clusterMinTransportVersion = clusterState.getMinTransportVersion();
         this.aliasFilter = aliasFilter;
         this.results = resultConsumer;
         // register the release of the query consumer to free up the circuit breaker memory
@@ -609,7 +612,7 @@ public void sendSearchResponse(SearchResponseSections internalSearchResponse, At
     protected BytesReference buildSearchContextId(ShardSearchFailure[] failures) {
         var source = request.source();
         return source != null && source.pointInTimeBuilder() != null && source.pointInTimeBuilder().singleSession() == false
-            ? source.pointInTimeBuilder().getEncodedId()
+            ? SearchContextId.encode(results.getAtomicArray().asList(), aliasFilter, clusterMinTransportVersion, failures)
             : null;
     }
 

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

@@ -26,7 +26,7 @@ public static SearchContextId decodePITId(String id) throws IOException {
         return decodePITId(new BytesArray(Base64.getUrlDecoder().decode(id)));
     }
 
-    public static SearchContextId decodePITId(BytesReference id) throws IOException {
+    public static SearchContextId decodePITId(BytesReference id) {
         try (var in = id.streamInput()) {
             final TransportVersion version = TransportVersion.readVersion(in);
             in.setTransportVersion(version);

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

@@ -30,6 +30,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.TreeSet;
 import java.util.stream.Collectors;
@@ -39,7 +40,7 @@ public final class SearchContextId {
     private final Map<String, AliasFilter> aliasFilter;
     private final transient Set<ShardSearchContextId> contextIds;
 
-    SearchContextId(Map<ShardId, SearchContextIdForNode> shards, Map<String, AliasFilter> aliasFilter) {
+    public SearchContextId(Map<ShardId, SearchContextIdForNode> shards, Map<String, AliasFilter> aliasFilter) {
         this.shards = shards;
         this.aliasFilter = aliasFilter;
         this.contextIds = shards.values().stream().map(SearchContextIdForNode::getSearchContextId).collect(Collectors.toSet());
@@ -57,8 +58,8 @@ public boolean contains(ShardSearchContextId contextId) {
         return contextIds.contains(contextId);
     }
 
-    public static BytesReference encode(
-        List<SearchPhaseResult> searchPhaseResults,
+    public static <SPR extends SearchPhaseResult> BytesReference encode(
+        List<SPR> searchPhaseResults,
         Map<String, AliasFilter> aliasFilter,
         TransportVersion version,
         ShardSearchFailure[] shardFailures
@@ -142,4 +143,17 @@ public String[] getActualIndices() {
         }
         return indices.toArray(String[]::new);
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (o == null || getClass() != o.getClass()) return false;
+        SearchContextId that = (SearchContextId) o;
+        return Objects.equals(shards, that.shards) && Objects.equals(aliasFilter,
+                that.aliasFilter) && Objects.equals(contextIds, that.contextIds);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(shards, aliasFilter, contextIds);
+    }
 }

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

@@ -17,6 +17,7 @@
 import org.elasticsearch.search.internal.ShardSearchContextId;
 
 import java.io.IOException;
+import java.util.Objects;
 
 public final class SearchContextIdForNode implements Writeable {
     private final String node;
@@ -30,7 +31,7 @@ public final class SearchContextIdForNode implements Writeable {
      * @param node The target node where the search context ID is defined, or {@code null} if the shard is missing or unavailable.
      * @param searchContextId The {@link ShardSearchContextId}, or {@code null} if the shard is missing or unavailable.
      */
-    SearchContextIdForNode(@Nullable String clusterAlias, @Nullable String node, @Nullable ShardSearchContextId searchContextId) {
+    public SearchContextIdForNode(@Nullable String clusterAlias, @Nullable String node, @Nullable ShardSearchContextId searchContextId) {
         this.node = node;
         this.clusterAlias = clusterAlias;
         this.searchContextId = searchContextId;
@@ -103,4 +104,18 @@ public String toString() {
             + '\''
             + '}';
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (o == null || getClass() != o.getClass()) return false;
+        SearchContextIdForNode that = (SearchContextIdForNode) o;
+        return Objects.equals(node, that.node)
+            && Objects.equals(searchContextId, that.searchContextId)
+            && Objects.equals(clusterAlias, that.clusterAlias);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(node, searchContextId, clusterAlias);
+    }
 }

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

@@ -709,7 +709,7 @@ protected void innerOnResponse(SearchPhaseResult searchPhaseResult) {
                             } catch (Exception e) {
                                 setFailure(state, dataNodeLocalIdx, e);
                             } finally {
-                                doneFuture.onResponse(null);
+                               doneFuture.onResponse(null);
                             }
                         }
 

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

@@ -1233,6 +1233,7 @@ static List<SearchShardIterator> getRemoteShardsIteratorFromPointInTime(
                     // Otherwise, we add the shard iterator without a target node, allowing a partial search failure to
                     // be thrown when a search phase attempts to access it.
                     targetNodes.add(perNode.getNode());
+                    // TODO we will need to adapt something here as well I think
                     if (perNode.getSearchContextId().getSearcherId() != null) {
                         for (String node : group.allocatedNodes()) {
                             if (node.equals(perNode.getNode()) == false) {

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

@@ -541,7 +541,6 @@ public void beforeIndexShardCreated(ShardRouting routing, Settings indexSettings
 
     protected void putReaderContext(ReaderContext context) {
         final ReaderContext previous = activeReaders.put(context.readerContextId(), context);
-        logger.info("--> adding reader context [{}], current [{}]", context.readerContextId(), activeReaders.size());
         assert previous == null;
         // ensure that if we race against afterIndexRemoved, we remove the context from the active list.
         // this is important to ensure store can be cleaned up, in particular if the search is a scroll with a long timeout.
@@ -556,8 +555,7 @@ protected ReaderContext removeReaderContext(ReaderContextId id) {
         if (logger.isTraceEnabled()) {
             logger.trace("removing reader context [{}]", id);
         }
-        ReaderContext remove = activeReaders.remove(id);
-        return remove;
+        return activeReaders.remove(id);
     }
 
     @Override
@@ -1261,6 +1259,7 @@ final ReaderContext createOrGetReaderContext(ShardSearchRequest request) {
                 if (searcherId == null) {
                     throw e;
                 }
+                // TODO this retries on context with same searcher id, currently offered in FrozenEngine and ReadonlyEngine
                 final IndexService indexService = indicesService.indexServiceSafe(request.shardId().getIndex());
                 final IndexShard shard = indexService.getShard(request.shardId().id());
                 final Engine.SearcherSupplier searcherSupplier = shard.acquireSearcherSupplier();