[ML] Reset retryable index requests after failures (#109320) (#109396)

Fixes the `autoGeneratedTimestamp should not be set externally` error

Author: davidkyle

docs/changelog/109320.yaml

@@ -0,0 +1,5 @@
+pr: 109320
+summary: Reset retryable index requests after failures
+area: Machine Learning
+type: bug
+issues: []

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/utils/persistence/ResultsPersisterService.java

@@ -325,7 +325,7 @@ private static boolean isIrrecoverable(Exception ex) {
     }
 
     @SuppressWarnings("NonAtomicOperationOnVolatileField")
-    private static class BulkRequestRewriter {
+    static class BulkRequestRewriter {
         private volatile BulkRequest bulkRequest;
 
         BulkRequestRewriter(BulkRequest initialRequest) {
@@ -533,7 +533,7 @@ public void cancel(Exception e) {
         }
     }
 
-    private static BulkRequest buildNewRequestFromFailures(BulkRequest bulkRequest, BulkResponse bulkResponse) {
+    static BulkRequest buildNewRequestFromFailures(BulkRequest bulkRequest, BulkResponse bulkResponse) {
         // If we failed, lets set the bulkRequest to be a collection of the failed requests
         BulkRequest bulkRequestOfFailures = new BulkRequest();
         Set<String> failedDocIds = Arrays.stream(bulkResponse.getItems())
@@ -542,6 +542,9 @@ private static BulkRequest buildNewRequestFromFailures(BulkRequest bulkRequest,
             .collect(Collectors.toSet());
         bulkRequest.requests().forEach(docWriteRequest -> {
             if (failedDocIds.contains(docWriteRequest.id())) {
+                if (docWriteRequest instanceof IndexRequest ir) {
+                    ir.reset();
+                }
                 bulkRequestOfFailures.add(docWriteRequest);
             }
         });

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/utils/persistence/ResultsPersisterServiceTests.java

@@ -60,6 +60,7 @@
 import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.nullValue;
 import static org.mockito.ArgumentMatchers.any;
@@ -375,6 +376,34 @@ public void testBulkRequestRetriesMsgHandlerIsCalled() {
         assertThat(lastMessage.get(), containsString("failed to index after [1] attempts. Will attempt again"));
     }
 
+    public void testBuildNewRequestFromFailures_resetsId() {
+        var bulkRequest = new BulkRequest();
+        var indexRequestAutoGeneratedId = new IndexRequest("index-foo");
+        indexRequestAutoGeneratedId.autoGenerateId();
+        var autoGenId = indexRequestAutoGeneratedId.id();
+        var plainIndexRequest = new IndexRequest("index-foo2").id("id-set");
+
+        bulkRequest.add(indexRequestAutoGeneratedId);
+        bulkRequest.add(plainIndexRequest);
+
+        var bulkResponse = mock(BulkResponse.class);
+
+        var failed = mock(BulkItemResponse.class);
+        when(failed.isFailed()).thenReturn(Boolean.TRUE);
+        when(failed.getId()).thenReturn(autoGenId);
+
+        var sucessful = mock(BulkItemResponse.class);
+        when(sucessful.isFailed()).thenReturn(Boolean.FALSE);
+
+        when(bulkResponse.getItems()).thenReturn(new BulkItemResponse[] { failed, sucessful });
+
+        var modifiedRequestForRetry = ResultsPersisterService.buildNewRequestFromFailures(bulkRequest, bulkResponse);
+        assertThat(modifiedRequestForRetry.requests(), hasSize(1)); // only the failed item is in the new request
+        assertThat(modifiedRequestForRetry.requests().get(0), instanceOf(IndexRequest.class));
+        var ir = (IndexRequest) modifiedRequestForRetry.requests().get(0);
+        assertEquals(ir.getAutoGeneratedTimestamp(), -1L); // failed request was reset
+    }
+
     private static <Response> Stubber doAnswerWithResponses(Response response1, Response response2) {
         return doAnswer(withResponse(response1)).doAnswer(withResponse(response2));
     }