Only retry bulk request items that failed

janusgraph-es/src/main/java/org/janusgraph/diskstorage/es/rest/RestElasticSearchClient.java

@@ -41,12 +41,15 @@
 import org.janusgraph.diskstorage.es.mapping.TypelessIndexMappings;
 import org.janusgraph.diskstorage.es.script.ESScriptResponse;
 import org.javatuples.Pair;
+import org.javatuples.Triplet;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -380,8 +383,7 @@ public void clearStore(String indexName, String storeName) throws IOException {
         }
     }
 
-    @Override
-    public void bulkRequest(List<ElasticSearchMutation> requests, String ingestPipeline) throws IOException {
+    private Pair<String, byte[]> buildBulkRequestInput(List<ElasticSearchMutation> requests, String ingestPipeline) throws IOException {
         final ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
         for (final ElasticSearchMutation request : requests) {
             Map<String, Object> requestData = new HashMap<>();
@@ -416,24 +418,49 @@ public void bulkRequest(List<ElasticSearchMutation> requests, String ingestPipel
             APPEND_OP.apply(builder).append("refresh=").append(bulkRefresh);
         }
         builder.insert(0, REQUEST_SEPARATOR + "_bulk");
+        return Pair.with(builder.toString(), outputStream.toByteArray());
+    }
+
+    private List<Triplet<Object, Integer, ElasticSearchMutation>> pairErrorsWithSubmittedMutation(
+        //Bulk API is documented to return bulk item responses in the same order of submission
+        //https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html#bulk-api-response-body
+        //As such we only need to retry elements that failed
+        final List<Map<String, RestBulkResponse.RestBulkItemResponse>> bulkResponseItems,
+        final List<ElasticSearchMutation> submittedBulkRequestItems) {
+        final List<Triplet<Object, Integer, ElasticSearchMutation>> errors = new ArrayList<>(bulkResponseItems.size());
+        for (int itemIndex = 0; itemIndex < bulkResponseItems.size(); itemIndex++) {
+            Collection<RestBulkResponse.RestBulkItemResponse> bulkResponseItem = bulkResponseItems.get(itemIndex).values();
+            if (bulkResponseItem.size() > 1) {
+                throw new IllegalStateException("There should only be a single item per bulk reponse item entry");
+            }
+            RestBulkResponse.RestBulkItemResponse item = bulkResponseItem.iterator().next();
+            if (item.getError() != null && item.getStatus() != 404) {
+                errors.add(Triplet.with(item.getError(), item.getStatus(), submittedBulkRequestItems.get(itemIndex)));
+            }
+        }
+        return errors;
+    }
 
+    @Override
+    public void bulkRequest(List<ElasticSearchMutation> requests, String ingestPipeline) throws IOException {
+        List<ElasticSearchMutation> requestsToSend = requests;
         int retryCount = 0;
         while (true) {
-            final Response response = performRequest(REQUEST_TYPE_POST, builder.toString(), outputStream.toByteArray());
+            final Pair<String, byte[]> bulkRequestInput = buildBulkRequestInput(requestsToSend, ingestPipeline);
+            final Response response = performRequest(REQUEST_TYPE_POST, bulkRequestInput.getValue0(), bulkRequestInput.getValue1());
             try (final InputStream inputStream = response.getEntity().getContent()) {
                 final RestBulkResponse bulkResponse = mapper.readValue(inputStream, RestBulkResponse.class);
-                final List<Pair<Object, Integer>> errors = bulkResponse.getItems().stream()
-                    .flatMap(item -> item.values().stream())
-                    .filter(item -> item.getError() != null && item.getStatus() != 404)
-                    .map(item -> Pair.with(item.getError(), item.getStatus())).collect(Collectors.toList());
-                if (!errors.isEmpty()) {
+                List<Triplet<Object, Integer, ElasticSearchMutation>> bulkItemsThatFailed = pairErrorsWithSubmittedMutation(bulkResponse.getItems(), requestsToSend);
+                if (!bulkItemsThatFailed.isEmpty()) {
                     //Only retry the bulk request if *all* the bulk response item error codes are retry error codes
-                    final Set<Integer> errorCodes = errors.stream().map(Pair::getValue1).collect(Collectors.toSet());
+                    final Set<Integer> errorCodes = bulkItemsThatFailed.stream().map(Triplet::getValue1).collect(Collectors.toSet());
                     if (retryCount < retryAttemptLimit && retryOnErrorCodes.containsAll(errorCodes)) {
+                        //Build up the next request batch, of only the failed mutations
+                        requestsToSend = bulkItemsThatFailed.stream().map(Triplet::getValue2).collect(Collectors.toList());
                         performRetryWait(retryCount);
                         retryCount++;
                     } else {
-                        final List<Object> errorItems = errors.stream().map(Pair::getValue0).collect(Collectors.toList());
+                        final List<Object> errorItems = bulkItemsThatFailed.stream().map(Triplet::getValue0).collect(Collectors.toList());
                         errorItems.forEach(error -> log.error("Failed to execute ES query: {}", error));
                         throw new IOException("Failure(s) in Elasticsearch bulk request: " + errorItems);
                     }