Adding support for allow_partial_search_results in PIT (elastic#111516)

Author: pmpailis

docs/changelog/111516.yaml

@@ -0,0 +1,5 @@
+pr: 111516
+summary: Adding support for `allow_partial_search_results` in PIT
+area: Search
+type: enhancement
+issues: []

docs/reference/search/point-in-time-api.asciidoc

@@ -78,6 +78,44 @@ IMPORTANT: The open point in time request and each subsequent search request can
 return different `id`; thus always use the most recently received `id` for the
 next search request.
 
+In addition to the `keep_alive` parameter, the `allow_partial_search_results` parameter
+can also be defined.
+This parameter determines whether the <<point-in-time-api, point in time (PIT)>>
+should tolerate unavailable shards or <<shard-failures, shard failures>> when
+initially creating the PIT.
+If set to true, the PIT will be created with the available shards, along with a
+reference to any missing ones.
+If set to false, the operation will fail if any shard is unavailable.
+The default value is false.
+
+The PIT response includes a summary of the total number of shards, as well as the number
+of successful shards when creating the PIT.
+
+[source,console]
+--------------------------------------------------
+POST /my-index-000001/_pit?keep_alive=1m&allow_partial_search_results=true
+--------------------------------------------------
+// TEST[setup:my_index]
+
+[source,js]
+--------------------------------------------------
+{
+  "id": "46ToAwMDaWR5BXV1aWQyKwZub2RlXzMAAAAAAAAAACoBYwADaWR4BXV1aWQxAgZub2RlXzEAAAAAAAAAAAEBYQADaWR5BXV1aWQyKgZub2RlXzIAAAAAAAAAAAwBYgACBXV1aWQyAAAFdXVpZDEAAQltYXRjaF9hbGw_gAAAAA=",
+  "_shards": {
+    "total": 10,
+    "successful": 10,
+    "skipped": 0,
+    "failed": 0
+  }
+}
+--------------------------------------------------
+// NOTCONSOLE
+
+When a PIT that contains shard failures is used in a search request, the missing are
+always reported in the search response as a NoShardAvailableActionException exception.
+To get rid of these exceptions, a new PIT needs to be created so that shards missing
+from the previous PIT can be handled, assuming they become available in the meantime.
+
 [[point-in-time-keep-alive]]
 ==== Keeping point in time alive
 The `keep_alive` parameter, which is passed to a open point in time request and

docs/reference/search/search-your-data/paginate-search-results.asciidoc

@@ -106,9 +106,9 @@ The search response includes an array of `sort` values for each hit:
         "_id" : "654322",
         "_score" : null,
         "_source" : ...,
-        "sort" : [               
+        "sort" : [
           1463538855,
-          "654322"                           
+          "654322"
         ]
       },
       {
@@ -118,7 +118,7 @@ The search response includes an array of `sort` values for each hit:
         "_source" : ...,
         "sort" : [                                <1>
           1463538857,
-          "654323"                        
+          "654323"
         ]
       }
     ]
@@ -150,7 +150,7 @@ GET twitter/_search
 --------------------------------------------------
 //TEST[continued]
 
-Repeat this process by updating the `search_after` array every time you retrieve a 
+Repeat this process by updating the `search_after` array every time you retrieve a
 new page of results. If a <<near-real-time,refresh>> occurs between these requests,
 the order of your results may change, causing inconsistent results across pages. To
 prevent this, you can create a <<point-in-time-api,point in time (PIT)>> to
@@ -167,10 +167,12 @@ The API returns a PIT ID.
 [source,console-result]
 ----
 {
-  "id": "46ToAwMDaWR5BXV1aWQyKwZub2RlXzMAAAAAAAAAACoBYwADaWR4BXV1aWQxAgZub2RlXzEAAAAAAAAAAAEBYQADaWR5BXV1aWQyKgZub2RlXzIAAAAAAAAAAAwBYgACBXV1aWQyAAAFdXVpZDEAAQltYXRjaF9hbGw_gAAAAA=="
+  "id": "46ToAwMDaWR5BXV1aWQyKwZub2RlXzMAAAAAAAAAACoBYwADaWR4BXV1aWQxAgZub2RlXzEAAAAAAAAAAAEBYQADaWR5BXV1aWQyKgZub2RlXzIAAAAAAAAAAAwBYgACBXV1aWQyAAAFdXVpZDEAAQltYXRjaF9hbGw_gAAAAA==",
+  "_shards": ...
 }
 ----
 // TESTRESPONSE[s/"id": "46ToAwMDaWR5BXV1aWQyKwZub2RlXzMAAAAAAAAAACoBYwADaWR4BXV1aWQxAgZub2RlXzEAAAAAAAAAAAEBYQADaWR5BXV1aWQyKgZub2RlXzIAAAAAAAAAAAwBYgACBXV1aWQyAAAFdXVpZDEAAQltYXRjaF9hbGw_gAAAAA=="/"id": $body.id/]
+// TESTRESPONSE[s/"_shards": \.\.\./"_shards": "$body._shards"/]
 
 To get the first page of results, submit a search request with a `sort`
 argument. If using a PIT, specify the PIT ID in the `pit.id` parameter and omit

server/src/internalClusterTest/java/org/elasticsearch/action/search/PointInTimeIT.java

@@ -195,6 +195,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ESQL_PROFILE_SLEEPS = def(8_725_00_0);
     public static final TransportVersion ZDT_NANOS_SUPPORT = def(8_726_00_0);
     public static final TransportVersion LTR_SERVERLESS_RELEASE = def(8_727_00_0);
+    public static final TransportVersion ALLOW_PARTIAL_SEARCH_RESULTS_IN_PIT = def(8_728_00_0);
     /*
      * STOP! READ THIS FIRST! No, really,
      *        ____ _____ ___  ____  _        ____  _____    _    ____    _____ _   _ ___ ____    _____ ___ ____  ____ _____ _

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -707,7 +707,7 @@ public void sendSearchResponse(SearchResponseSections internalSearchResponse, At
             final String scrollId = request.scroll() != null ? TransportSearchHelper.buildScrollId(queryResults) : null;
             final BytesReference searchContextId;
             if (buildPointInTimeFromSearchResults()) {
-                searchContextId = SearchContextId.encode(queryResults.asList(), aliasFilter, minTransportVersion);
+                searchContextId = SearchContextId.encode(queryResults.asList(), aliasFilter, minTransportVersion, failures);
             } else {
                 if (request.source() != null
                     && request.source().pointInTimeBuilder() != null

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

@@ -166,6 +166,10 @@ public static void closeContexts(
             final var successes = new AtomicInteger();
             try (RefCountingRunnable refs = new RefCountingRunnable(() -> l.onResponse(successes.get()))) {
                 for (SearchContextIdForNode contextId : contextIds) {
+                    if (contextId.getNode() == null) {
+                        // the shard was missing when creating the PIT, ignore.
+                        continue;
+                    }
                     final DiscoveryNode node = nodeLookup.apply(contextId.getClusterAlias(), contextId.getNode());
                     if (node != null) {
                         try {

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

@@ -41,6 +41,8 @@ public final class OpenPointInTimeRequest extends ActionRequest implements Indic
 
     private QueryBuilder indexFilter;
 
+    private boolean allowPartialSearchResults = false;
+
     public static final IndicesOptions DEFAULT_INDICES_OPTIONS = SearchRequest.DEFAULT_INDICES_OPTIONS;
 
     public OpenPointInTimeRequest(String... indices) {
@@ -60,6 +62,9 @@ public OpenPointInTimeRequest(StreamInput in) throws IOException {
         if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_12_0)) {
             this.indexFilter = in.readOptionalNamedWriteable(QueryBuilder.class);
         }
+        if (in.getTransportVersion().onOrAfter(TransportVersions.ALLOW_PARTIAL_SEARCH_RESULTS_IN_PIT)) {
+            this.allowPartialSearchResults = in.readBoolean();
+        }
     }
 
     @Override
@@ -76,6 +81,11 @@ public void writeTo(StreamOutput out) throws IOException {
         if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_12_0)) {
             out.writeOptionalWriteable(indexFilter);
         }
+        if (out.getTransportVersion().onOrAfter(TransportVersions.ALLOW_PARTIAL_SEARCH_RESULTS_IN_PIT)) {
+            out.writeBoolean(allowPartialSearchResults);
+        } else if (allowPartialSearchResults) {
+            throw new IOException("[allow_partial_search_results] is not supported on nodes with version " + out.getTransportVersion());
+        }
     }
 
     @Override
@@ -180,6 +190,15 @@ public boolean includeDataStreams() {
         return true;
     }
 
+    public boolean allowPartialSearchResults() {
+        return allowPartialSearchResults;
+    }
+
+    public OpenPointInTimeRequest allowPartialSearchResults(boolean allowPartialSearchResults) {
+        this.allowPartialSearchResults = allowPartialSearchResults;
+        return this;
+    }
+
     @Override
     public String getDescription() {
         return "open search context: indices [" + String.join(",", indices) + "] keep_alive [" + keepAlive + "]";
@@ -200,6 +219,8 @@ public String toString() {
             + ", preference='"
             + preference
             + '\''
+            + ", allowPartialSearchResults="
+            + allowPartialSearchResults
             + '}';
     }
 
@@ -218,12 +239,13 @@ public boolean equals(Object o) {
             && indicesOptions.equals(that.indicesOptions)
             && keepAlive.equals(that.keepAlive)
             && Objects.equals(routing, that.routing)
-            && Objects.equals(preference, that.preference);
+            && Objects.equals(preference, that.preference)
+            && Objects.equals(allowPartialSearchResults, that.allowPartialSearchResults);
     }
 
     @Override
     public int hashCode() {
-        int result = Objects.hash(indicesOptions, keepAlive, maxConcurrentShardRequests, routing, preference);
+        int result = Objects.hash(indicesOptions, keepAlive, maxConcurrentShardRequests, routing, preference, allowPartialSearchResults);
         result = 31 * result + Arrays.hashCode(indices);
         return result;
     }

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

@@ -8,6 +8,7 @@
 
 package org.elasticsearch.action.search;
 
+import org.elasticsearch.TransportVersions;
 import org.elasticsearch.action.ActionResponse;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.io.stream.StreamOutput;
@@ -18,22 +19,46 @@
 import java.util.Base64;
 import java.util.Objects;
 
+import static org.elasticsearch.rest.action.RestActions.buildBroadcastShardsHeader;
+
 public final class OpenPointInTimeResponse extends ActionResponse implements ToXContentObject {
     private final BytesReference pointInTimeId;
 
-    public OpenPointInTimeResponse(BytesReference pointInTimeId) {
+    private final int totalShards;
+    private final int successfulShards;
+    private final int failedShards;
+    private final int skippedShards;
+
+    public OpenPointInTimeResponse(
+        BytesReference pointInTimeId,
+        int totalShards,
+        int successfulShards,
+        int failedShards,
+        int skippedShards
+    ) {
         this.pointInTimeId = Objects.requireNonNull(pointInTimeId, "Point in time parameter must be not null");
+        this.totalShards = totalShards;
+        this.successfulShards = successfulShards;
+        this.failedShards = failedShards;
+        this.skippedShards = skippedShards;
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeBytesReference(pointInTimeId);
+        if (out.getTransportVersion().onOrAfter(TransportVersions.ALLOW_PARTIAL_SEARCH_RESULTS_IN_PIT)) {
+            out.writeVInt(totalShards);
+            out.writeVInt(successfulShards);
+            out.writeVInt(failedShards);
+            out.writeVInt(skippedShards);
+        }
     }
 
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
         builder.field("id", Base64.getUrlEncoder().encodeToString(BytesReference.toBytes(pointInTimeId)));
+        buildBroadcastShardsHeader(builder, params, totalShards, successfulShards, failedShards, skippedShards, null);
         builder.endObject();
         return builder;
     }
@@ -42,4 +67,19 @@ public BytesReference getPointInTimeId() {
         return pointInTimeId;
     }
 
+    public int getTotalShards() {
+        return totalShards;
+    }
+
+    public int getSuccessfulShards() {
+        return successfulShards;
+    }
+
+    public int getFailedShards() {
+        return failedShards;
+    }
+
+    public int getSkippedShards() {
+        return skippedShards;
+    }
 }

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

@@ -47,6 +47,7 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
         openRequest.routing(request.param("routing"));
         openRequest.preference(request.param("preference"));
         openRequest.keepAlive(TimeValue.parseTimeValue(request.param("keep_alive"), null, "keep_alive"));
+        openRequest.allowPartialSearchResults(request.paramAsBoolean("allow_partial_search_results", false));
         if (request.hasParam("max_concurrent_shard_requests")) {
             final int maxConcurrentShardRequests = request.paramAsInt(
                 "max_concurrent_shard_requests",