ResumeInfo never nullable

Author: szybia

modules/reindex-management/src/main/java/org/elasticsearch/reindex/management/TransportListReindexAction.java

@@ -92,28 +92,26 @@ protected void taskOperation(CancellableTask actionTask, ListReindexRequest requ
         assert task instanceof BulkByScrollTask : "task should be a BulkByScrollTask";
         TaskInfo info = task.taskInfo(clusterService.localNode().getId(), request.getDetailed());
         if (task instanceof final BulkByScrollTask bbs) {
-            final ResumeInfo.RelocationOrigin origin = bbs.getRelocationOrigin().orElse(null);
-            if (origin != null) {
-                final TaskId originalId = origin.originalTaskId();
-                final long originalStartMillis = origin.originalStartTimeMillis();
-                final long adjustedRunningTimeNanos = info.runningTimeNanos() + TimeUnit.MILLISECONDS.toNanos(
-                    info.startTime() - originalStartMillis
-                );
-                info = new TaskInfo(
-                    originalId,
-                    info.type(),
-                    originalId.getNodeId(),
-                    info.action(),
-                    info.description(),
-                    info.status(),
-                    originalStartMillis,
-                    adjustedRunningTimeNanos,
-                    info.cancellable(),
-                    info.cancelled(),
-                    info.parentTaskId(),
-                    info.headers()
-                );
-            }
+            final ResumeInfo.RelocationOrigin origin = bbs.getRelocationOrigin();
+            final TaskId originalId = origin.originalTaskId();
+            final long originalStartMillis = origin.originalStartTimeMillis();
+            final long adjustedRunningTimeNanos = info.runningTimeNanos() + TimeUnit.MILLISECONDS.toNanos(
+                info.startTime() - originalStartMillis
+            );
+            info = new TaskInfo(
+                originalId,
+                info.type(),
+                originalId.getNodeId(),
+                info.action(),
+                info.description(),
+                info.status(),
+                originalStartMillis,
+                adjustedRunningTimeNanos,
+                info.cancellable(),
+                info.cancelled(),
+                info.parentTaskId(),
+                info.headers()
+            );
         }
         listener.onResponse(info);
     }

modules/reindex/src/main/java/org/elasticsearch/reindex/AbstractAsyncBulkByScrollAction.java

@@ -570,7 +570,7 @@ void notifyDone(long thisBatchStartTimeNS, ScrollConsumableHitsResponse asyncRes
                     worker.getStatus(),
                     remoteVersion
                 );
-                final ResumeInfo resumeInfo = new ResumeInfo(workerResumeInfo, null);
+                final ResumeInfo resumeInfo = new ResumeInfo(task.getRelocationOrigin(), workerResumeInfo, null);
                 // This response is a local carrier for resumeInfo — for higher-level code to handle relocation and then discard.
                 // However, status must be accurate for sliced tasks only, the leader state stores this response and derives
                 // its own combined status from it to serialize to .tasks index.

modules/reindex/src/main/java/org/elasticsearch/reindex/Reindexer.java

@@ -143,12 +143,22 @@ public class Reindexer {
 
     public void initTask(BulkByScrollTask task, ReindexRequest request, ActionListener<Void> listener) {
         final ActionListener<Void> initListener = listener.delegateFailure((l, v) -> {
+            initRelocationOrigin(task, request);
             initTaskForRelocationIfEnabled(task, request);
             l.onResponse(v);
         });
         BulkByPaginatedSearchParallelizationHelper.initTaskState(task, request, client, initListener);
     }
 
+    private void initRelocationOrigin(BulkByScrollTask task, ReindexRequest request) {
+        if (request.getRelocationOrigin() != null) {
+            return;
+        }
+        var selfOrigin = new ResumeInfo.RelocationOrigin(new TaskId(clusterService.localNode().getId(), task.getId()), task.getStartTime());
+        task.initSelfOrigin(selfOrigin);
+        request.setRelocationOrigin(selfOrigin);
+    }
+
     public void execute(BulkByScrollTask task, ReindexRequest request, Client bulkClient, ActionListener<BulkByScrollResponse> listener) {
         // todo(szy/sam): handle sliced and non-sliced
         // todo(szy/sam): bug, we send System::nanoTime across JVMs
@@ -373,12 +383,7 @@ ActionListener<BulkByScrollResponse> listenerWithRelocations(
                 return;
             }
             final ResumeInfo currentResumeInfo = response.getTaskResumeInfo().get();
-            final ResumeInfo.RelocationOrigin origin = request.getResumeInfo()
-                .map(r -> Objects.requireNonNull(r.relocationOrigin(), "relocation origin should be set if resume info is present"))
-                .orElseGet(
-                    () -> new ResumeInfo.RelocationOrigin(new TaskId(clusterService.localNode().getId(), task.getId()), task.getStartTime())
-                );
-            request.setResumeInfo(new ResumeInfo(origin, currentResumeInfo.worker(), currentResumeInfo.slices()));
+            request.setResumeInfo(new ResumeInfo(task.getRelocationOrigin(), currentResumeInfo.worker(), currentResumeInfo.slices()));
             final ResumeBulkByScrollRequest resumeRequest = new ResumeBulkByScrollRequest(request);
             final ActionListener<ResumeBulkByScrollResponse> relocationListener = ActionListener.wrap(resp -> {
                 final var relocatedException = new TaskRelocatedException();

server/src/main/java/org/elasticsearch/index/reindex/AbstractBulkByScrollRequest.java

@@ -116,6 +116,14 @@ public abstract class AbstractBulkByScrollRequest<Self extends AbstractBulkByScr
     @Nullable
     private ResumeInfo resumeInfo;
 
+    /**
+     * Identity of the original task, propagated across relocations and slices so the user-facing task ID and start time are preserved.
+     * Transient — not serialized; flows through {@link ResumeInfo} on the wire, then synced here via {@link #setResumeInfo} or set
+     * directly by the init path for first-run tasks.
+     */
+    @Nullable
+    private ResumeInfo.RelocationOrigin relocationOrigin;
+
     public AbstractBulkByScrollRequest(StreamInput in) throws IOException {
         super(in);
         searchRequest = new SearchRequest(in);
@@ -134,7 +142,13 @@ public AbstractBulkByScrollRequest(StreamInput in) throws IOException {
             eligibleForRelocationOnShutdown = in.readBoolean();
         }
         if (in.getTransportVersion().supports(REINDEX_RELOCATION_RESUME)) {
-            resumeInfo = in.readOptionalWriteable(ResumeInfo::new);
+            var ri = in.readOptionalWriteable(ResumeInfo::new);
+            this.resumeInfo = ri;
+            if (ri != null) {
+                if (resumeInfo.relocationOrigin() != null) {
+                    this.relocationOrigin = resumeInfo.relocationOrigin();
+                }
+            }
         }
     }
 
@@ -437,9 +451,13 @@ public int getSlices() {
 
     /**
      * Sets resumption data to continue from a previously-acquired scroll ID.
+     * Also syncs {@link #relocationOrigin} from the resume info when present.
      */
     public Self setResumeInfo(ResumeInfo resumeInfo) {
         this.resumeInfo = Objects.requireNonNull(resumeInfo);
+        if (resumeInfo.relocationOrigin() != null) {
+            this.relocationOrigin = resumeInfo.relocationOrigin();
+        }
         return self();
     }
 
@@ -450,6 +468,16 @@ public Optional<ResumeInfo> getResumeInfo() {
         return Optional.ofNullable(resumeInfo);
     }
 
+    public Self setRelocationOrigin(ResumeInfo.RelocationOrigin relocationOrigin) {
+        this.relocationOrigin = Objects.requireNonNull(relocationOrigin);
+        return self();
+    }
+
+    @Nullable
+    public ResumeInfo.RelocationOrigin getRelocationOrigin() {
+        return relocationOrigin;
+    }
+
     /**
      * Build a new request for a slice of the parent request.
      */
@@ -475,12 +503,15 @@ protected Self doForSlice(Self request, TaskId slicingTask, int totalSlices) {
             .setRequestsPerSecond(requestsPerSecond / totalSlices)
             // Sub requests don't have workers
             .setSlices(1);
+        if (this.relocationOrigin != null) {
+            request.setRelocationOrigin(this.relocationOrigin);
+        }
         // Copy resume info for the slice from leader to the slice request
         if (this.getResumeInfo().isPresent()) {
             ResumeInfo resumeInfo = this.getResumeInfo().get();
             int sliceId = request.getSearchRequest().source().slice().getId();
             if (resumeInfo.isSliceCompleted(sliceId) == false) {
-                request.setResumeInfo(new ResumeInfo(resumeInfo.getSlice(sliceId).get().resumeInfo(), null));
+                request.setResumeInfo(new ResumeInfo(resumeInfo.relocationOrigin(), resumeInfo.getSlice(sliceId).get().resumeInfo(), null));
             }
         }
 
@@ -497,16 +528,7 @@ protected Self doForSlice(Self request, TaskId slicingTask, int totalSlices) {
 
     @Override
     public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
-        return new BulkByScrollTask(
-            id,
-            type,
-            action,
-            getDescription(),
-            parentTaskId,
-            headers,
-            eligibleForRelocationOnShutdown,
-            resumeInfo == null ? null : resumeInfo.relocationOrigin()
-        );
+        return new BulkByScrollTask(id, type, action, getDescription(), parentTaskId, headers, eligibleForRelocationOnShutdown, relocationOrigin);
     }
 
     @Override

server/src/main/java/org/elasticsearch/index/reindex/BulkByScrollTask.java

@@ -33,7 +33,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
@@ -58,7 +57,7 @@ public class BulkByScrollTask extends CancellableTask {
 
     private final boolean eligibleForRelocationOnShutdown;
     @Nullable
-    private final ResumeInfo.RelocationOrigin relocationOrigin;
+    private volatile ResumeInfo.RelocationOrigin relocationOrigin;
     private volatile LeaderBulkByScrollTaskState leaderState;
     private volatile WorkerBulkByScrollTaskState workerState;
     private volatile boolean relocationRequested = false;
@@ -227,9 +226,19 @@ public boolean isRelocationRequested() {
         return relocationRequested;
     }
 
-    /** Returns the relocation origin if this task is a relocated continuation. */
-    public Optional<ResumeInfo.RelocationOrigin> getRelocationOrigin() {
-        return Optional.ofNullable(relocationOrigin);
+    /**
+     * Sets the self-origin for a first-run (non-relocated) task. Must only be called once, and only when no origin was provided at
+     * construction time. After this call, {@link #getRelocationOrigin()} is guaranteed non-null.
+     */
+    public void initSelfOrigin(ResumeInfo.RelocationOrigin origin) {
+        assert this.relocationOrigin == null : "initSelfOrigin called but relocationOrigin is already set";
+        this.relocationOrigin = Objects.requireNonNull(origin);
+    }
+
+    /** Returns the relocation origin for this task. Always non-null after task initialisation. */
+    public ResumeInfo.RelocationOrigin getRelocationOrigin() {
+        assert relocationOrigin != null : "relocationOrigin has not been initialised";
+        return relocationOrigin;
     }
 
     /**

server/src/main/java/org/elasticsearch/index/reindex/LeaderBulkByScrollTaskState.java

@@ -174,7 +174,7 @@ private Optional<BulkByScrollResponse> relocationResponseIfNeeded() {
         if (allJobsCompletedThereforeNoNeedForRelocation) {
             return Optional.empty();
         }
-        final var resumeInfo = new ResumeInfo(null, sliceResumeInfoMap);
+        final var resumeInfo = new ResumeInfo(task.getRelocationOrigin(), null, sliceResumeInfoMap);
         // this response is a local carrier for resumeInfo only — for higher-level code to handle relocation and then discard.
         // the status for the task that's serialized into the .tasks index is taken from the leader state.
         return Optional.of(