Retry bulk request if all errors are configured retry error codes

Signed-off-by: Allan Clements <[email protected]>

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

@@ -39,8 +39,8 @@
 import org.janusgraph.diskstorage.es.mapping.IndexMapping;
 import org.janusgraph.diskstorage.es.mapping.TypedIndexMappings;
 import org.janusgraph.diskstorage.es.mapping.TypelessIndexMappings;
-import org.janusgraph.diskstorage.es.rest.RestBulkResponse.RestBulkItemResponse;
 import org.janusgraph.diskstorage.es.script.ESScriptResponse;
+import org.javatuples.Pair;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -417,16 +417,29 @@ public void bulkRequest(List<ElasticSearchMutation> requests, String ingestPipel
         }
         builder.insert(0, REQUEST_SEPARATOR + "_bulk");
 
-        final Response response = performRequest(REQUEST_TYPE_POST, builder.toString(), outputStream.toByteArray());
-        try (final InputStream inputStream = response.getEntity().getContent()) {
-            final RestBulkResponse bulkResponse = mapper.readValue(inputStream, RestBulkResponse.class);
-            final List<Object> errors = bulkResponse.getItems().stream()
-                .flatMap(item -> item.values().stream())
-                .filter(item -> item.getError() != null && item.getStatus() != 404)
-                .map(RestBulkItemResponse::getError).collect(Collectors.toList());
-            if (!errors.isEmpty()) {
-                errors.forEach(error -> log.error("Failed to execute ES query: {}", error));
-                throw new IOException("Failure(s) in Elasticsearch bulk request: " + errors);
+        int retryCount = 0;
+        while (true) {
+            final Response response = performRequest(REQUEST_TYPE_POST, builder.toString(), outputStream.toByteArray());
+            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()) {
+                    //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());
+                    if (retryCount < retryAttemptLimit && retryOnErrorCodes.containsAll(errorCodes)) {
+                        performRetryWait(retryCount);
+                        retryCount++;
+                    } else {
+                        errors.forEach(error -> log.error("Failed to execute ES query: {}", error.getValue0()));
+                        throw new IOException("Failure(s) in Elasticsearch bulk request: " + errors);
+                    }
+                } else {
+                    //The entire bulk request was successful, leave the loop
+                    break;
+                }
             }
         }
     }
@@ -571,19 +584,22 @@ private Response performRequestWithRetry(Request request) throws IOException {
                 if (!retryOnErrorCodes.contains(e.getResponse().getStatusLine().getStatusCode()) || retryCount >= retryAttemptLimit) {
                     throw e;
                 }
-                //Wait before trying again
-                long waitDurationMs = Math.min((long) (retryInitialWaitMs * Math.pow(10, retryCount)), retryMaxWaitMs);
-                log.warn("Retrying Elasticsearch request in {} ms. Attempt {} of {}", waitDurationMs, retryCount, retryAttemptLimit);
-                try {
-                    Thread.sleep(waitDurationMs);
-                } catch (InterruptedException interruptedException) {
-                    throw new RuntimeException(String.format("Thread interrupted while waiting for retry attempt %d of %d", retryCount, retryAttemptLimit), interruptedException);
-                }
+                performRetryWait(retryCount);
             }
             retryCount++;
         }
     }
 
+    private void performRetryWait(int retryCount) {
+        long waitDurationMs = Math.min((long) (retryInitialWaitMs * Math.pow(10, retryCount)), retryMaxWaitMs);
+        log.warn("Retrying Elasticsearch request in {} ms. Attempt {} of {}", waitDurationMs, retryCount, retryAttemptLimit);
+        try {
+            Thread.sleep(waitDurationMs);
+        } catch (InterruptedException interruptedException) {
+            throw new RuntimeException(String.format("Thread interrupted while waiting for retry attempt %d of %d", retryCount, retryAttemptLimit), interruptedException);
+        }
+    }
+
     private Response performRequest(Request request, byte[] requestData) throws IOException {
 
         final HttpEntity entity = requestData != null ? new ByteArrayEntity(requestData, ContentType.APPLICATION_JSON) : null;