From b60b0b9bbe5b9cf1b67af107dc6e44e4f2743500 Mon Sep 17 00:00:00 2001
From: Daniel Mitterdorfer <daniel.mitterdorfer@elastic.co>
Date: Thu, 26 Oct 2023 08:47:40 +0200
Subject: [PATCH] Handle missing locations

---
 .../get/TransportShardMultiGetAction.java     | 36 +++++++++++++------
 1 file changed, 26 insertions(+), 10 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/action/get/TransportShardMultiGetAction.java b/server/src/main/java/org/elasticsearch/action/get/TransportShardMultiGetAction.java
index 08a9314e53e6a..91808047397bf 100644
--- a/server/src/main/java/org/elasticsearch/action/get/TransportShardMultiGetAction.java
+++ b/server/src/main/java/org/elasticsearch/action/get/TransportShardMultiGetAction.java
@@ -51,7 +51,7 @@
 import static org.elasticsearch.core.Strings.format;
 
 public class TransportShardMultiGetAction extends TransportSingleShardAction<MultiGetShardRequest, MultiGetShardResponse> {
-    private static final boolean USE_BATCHED_GET  = Booleans.parseBoolean(System.getProperty("es.mget.batched", "false"));
+    private static final boolean USE_BATCHED_GET = Booleans.parseBoolean(System.getProperty("es.mget.batched", "false"));
 
     private static final String ACTION_NAME = MultiGetAction.NAME + "[shard]";
     public static final ActionType<MultiGetShardResponse> TYPE = new ActionType<>(ACTION_NAME, MultiGetShardResponse::new);
@@ -145,7 +145,7 @@ protected void asyncShardOperation(MultiGetShardRequest request, ShardId shardId
     @Override
     protected MultiGetShardResponse shardOperation(MultiGetShardRequest request, ShardId shardId) {
         MultiGetShardResponse response = new MultiGetShardResponse();
-        getAndAddToResponse(request, response, shardId);
+        getAndAddToResponse(request, response, shardId, false);
         return response;
     }
 
@@ -229,27 +229,37 @@ private void handleMultiGetOnUnpromotableShard(
 
     private MultiGetShardResponse handleLocalGets(MultiGetShardRequest request, MultiGetShardResponse response, ShardId shardId) {
         logger.trace("handling local gets for missing locations");
-        getAndAddToResponse(request, response, shardId);
+        getAndAddToResponse(request, response, shardId, true);
         return response;
     }
 
-    private void getAndAddToResponse(MultiGetShardRequest request, MultiGetShardResponse response, ShardId shardId) {
+    private void getAndAddToResponse(
+        MultiGetShardRequest request,
+        MultiGetShardResponse response,
+        ShardId shardId,
+        boolean expectMissingLocations
+    ) {
         long start = System.nanoTime();
         if (USE_BATCHED_GET) {
-            getAndAddToResponseBatched(request, response, shardId);
+            getAndAddToResponseBatched(request, response, shardId, expectMissingLocations);
             logger.info("Batched mget took=[" + (System.nanoTime() - start) / 1_000_000 + "] ms.");
         } else {
-            getAndAddToResponseSingle(request, response, shardId);
+            getAndAddToResponseSingle(request, response, shardId, expectMissingLocations);
             logger.info("Single mget took=[" + (System.nanoTime() - start) / 1_000_000 + "] ms.");
         }
     }
 
-    private void getAndAddToResponseSingle(MultiGetShardRequest request, MultiGetShardResponse response, ShardId shardId) {
+    private void getAndAddToResponseSingle(
+        MultiGetShardRequest request,
+        MultiGetShardResponse response,
+        ShardId shardId,
+        boolean expectMissingLocations
+    ) {
         var indexShard = getIndexShard(shardId);
         final int numRequests = request.locations.size();
         final GetAndFetchContext[] getAndFetchContexts = new GetAndFetchContext[numRequests];
         for (int location = 0; location < request.locations.size(); location++) {
-            if (response.responses.get(location) == null && response.failures.get(location) == null) {
+            if (expectMissingLocations == false || (response.responses.get(location) == null && response.failures.get(location) == null)) {
                 MultiGetRequest.Item item = request.items.get(location);
                 try {
                     GetResult getResult = indexShard.getService()
@@ -277,12 +287,18 @@ private void getAndAddToResponseSingle(MultiGetShardRequest request, MultiGetSha
             }
         }
     }
-    private void getAndAddToResponseBatched(MultiGetShardRequest request, MultiGetShardResponse response, ShardId shardId) {
+
+    private void getAndAddToResponseBatched(
+        MultiGetShardRequest request,
+        MultiGetShardResponse response,
+        ShardId shardId,
+        boolean expectMissingLocations
+    ) {
         var indexShard = getIndexShard(shardId);
         final int numRequests = request.locations.size();
         final GetAndFetchContext[] getAndFetchContexts = new GetAndFetchContext[numRequests];
         for (int location = 0; location < request.locations.size(); location++) {
-            if (response.responses.get(location) == null && response.failures.get(location) == null) {
+            if (expectMissingLocations == false || (response.responses.get(location) == null && response.failures.get(location) == null)) {
                 MultiGetRequest.Item item = request.items.get(location);
                 // try {
                 /*