TODO refactor: Use TaskId in TaskAwareRequest.createRequest

GalLalouche · GalLalouche · commit 0f9d2bd07671 · 2025-06-29T17:43:36.000+03:00
Use TaskId instead of passing a node name and ID individually.
diff --git a/modules/streams/src/main/java/org/elasticsearch/rest/streams/logs/LogsStreamsActivationToggleAction.java b/modules/streams/src/main/java/org/elasticsearch/rest/streams/logs/LogsStreamsActivationToggleAction.java
@@ -56,8 +56,8 @@ public boolean shouldEnable() {
         }
 
         @Override
-        public Task createTask(String localNodeId, long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
-            return new CancellableTask(id, type, action, "Logs streams activation toggle request", parentTaskId, headers);
+        public Task createTask(TaskId taskId, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+            return new CancellableTask(taskId.getId(), type, action, "Logs streams activation toggle request", parentTaskId, headers);
         }
     }
 }
diff --git a/plugins/discovery-ec2/src/test/java/org/elasticsearch/discovery/ec2/EC2RetriesTests.java b/plugins/discovery-ec2/src/test/java/org/elasticsearch/discovery/ec2/EC2RetriesTests.java
@@ -61,6 +61,8 @@ protected MockTransportService createTransportService() {
             ),
             threadPool,
             TransportService.NOOP_TRANSPORT_INTERCEPTOR,
+            null,
+            null,
             null
         );
     }
diff --git a/plugins/discovery-ec2/src/test/java/org/elasticsearch/discovery/ec2/Ec2DiscoveryTests.java b/plugins/discovery-ec2/src/test/java/org/elasticsearch/discovery/ec2/Ec2DiscoveryTests.java
@@ -78,7 +78,15 @@ public TransportAddress[] addressesFromString(String address) {
                 return new TransportAddress[] { poorMansDNS.getOrDefault(address, buildNewFakeTransportAddress()) };
             }
         };
-        return new MockTransportService(Settings.EMPTY, transport, threadPool, TransportService.NOOP_TRANSPORT_INTERCEPTOR, null);
+        return new MockTransportService(
+            Settings.EMPTY,
+            transport,
+            threadPool,
+            TransportService.NOOP_TRANSPORT_INTERCEPTOR,
+            null,
+            null,
+            null
+        );
     }
 
     protected List<TransportAddress> buildDynamicHosts(Settings nodeSettings, int nodes) {
diff --git a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java
@@ -685,14 +685,16 @@ private void construct(
 
         modules.bindToInstance(Tracer.class, telemetryProvider.getTracer());
 
+        assert nodeEnvironment.nodeId() != null : "node ID must be set before constructing the Node";
         TaskManager taskManager = new TaskManager(
             settings,
             threadPool,
             Stream.concat(
                 pluginsService.filterPlugins(ActionPlugin.class).flatMap(p -> p.getTaskHeaders().stream()),
                 Task.HEADERS_TO_COPY.stream()
             ).collect(Collectors.toSet()),
-            telemetryProvider.getTracer()
+            telemetryProvider.getTracer(),
+            nodeEnvironment.nodeId()
         );
 
         ClusterService clusterService = createClusterService(settingsModule, threadPool, taskManager);
@@ -1097,7 +1099,8 @@ public Map<String, String> queryFields() {
             localNodeFactory,
             settingsModule.getClusterSettings(),
             taskManager,
-            telemetryProvider.getTracer()
+            telemetryProvider.getTracer(),
+            nodeEnvironment.nodeId()
         );
         final SearchResponseMetrics searchResponseMetrics = new SearchResponseMetrics(telemetryProvider.getMeterRegistry());
         final SearchTransportService searchTransportService = new SearchTransportService(
diff --git a/server/src/main/java/org/elasticsearch/node/NodeServiceProvider.java b/server/src/main/java/org/elasticsearch/node/NodeServiceProvider.java
@@ -114,7 +114,8 @@ TransportService newTransportService(
         Function<BoundTransportAddress, DiscoveryNode> localNodeFactory,
         ClusterSettings clusterSettings,
         TaskManager taskManager,
-        Tracer tracer
+        Tracer tracer,
+        String initialNodeId
     ) {
         return new TransportService(settings, transport, threadPool, interceptor, localNodeFactory, clusterSettings, taskManager, tracer);
     }
diff --git a/server/src/main/java/org/elasticsearch/tasks/TaskAwareRequest.java b/server/src/main/java/org/elasticsearch/tasks/TaskAwareRequest.java
@@ -9,8 +9,6 @@
 
 package org.elasticsearch.tasks;
 
-import org.elasticsearch.core.Nullable;
-
 import java.util.Map;
 
 /**
@@ -57,18 +55,8 @@ default Task createTask(long id, String type, String action, TaskId parentTaskId
     /**
      * Returns the task object that should be used to keep track of the processing of the request, with an extra local node ID.
      */
-    // TODO remove the above overload, use only this one.
-    default Task createTask(
-        // TODO this is only nullable in tests, where the MockNode does not guarantee the localNodeId is set before calling this method. We
-        // We should fix the tests, and replace this and id with TaskId instead.
-        @Nullable String localNodeId,
-        long id,
-        String type,
-        String action,
-        TaskId parentTaskId,
-        Map<String, String> headers
-    ) {
-        return createTask(id, type, action, parentTaskId, headers);
+    default Task createTask(TaskId taskId, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+        return createTask(taskId.getId(), type, action, parentTaskId, headers);
     }
 
     /**
diff --git a/server/src/main/java/org/elasticsearch/tasks/TaskManager.java b/server/src/main/java/org/elasticsearch/tasks/TaskManager.java
@@ -21,6 +21,7 @@
 import org.elasticsearch.cluster.ClusterStateApplier;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.node.DiscoveryNodes;
+import org.elasticsearch.common.UUIDs;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
@@ -48,6 +49,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
@@ -80,6 +82,7 @@ public class TaskManager implements ClusterStateApplier {
 
     private TaskResultsService taskResultsService;
 
+    private final String initialNodeId;
     private DiscoveryNodes lastDiscoveryNodes = DiscoveryNodes.EMPTY_NODES;
 
     private final Tracer tracer;
@@ -96,10 +99,15 @@ public TaskManager(Settings settings, ThreadPool threadPool, Set<String> taskHea
     }
 
     public TaskManager(Settings settings, ThreadPool threadPool, Set<String> taskHeaders, Tracer tracer) {
+        this(settings, threadPool, taskHeaders, tracer, UUIDs.randomBase64UUID());
+    }
+
+    public TaskManager(Settings settings, ThreadPool threadPool, Set<String> taskHeaders, Tracer tracer, @Nullable String initialNodeId) {
         this.threadPool = threadPool;
         this.taskHeaders = Set.copyOf(taskHeaders);
         this.maxHeaderSize = SETTING_HTTP_MAX_HEADER_SIZE.get(settings);
         this.tracer = tracer;
+        this.initialNodeId = initialNodeId;
     }
 
     public void setTaskResultsService(TaskResultsService taskResultsService) {
@@ -141,9 +149,9 @@ public Task register(String type, String action, TaskAwareRequest request, boole
                 headers.put(key, httpHeader);
             }
         }
+        String localNodeId = Optional.ofNullable(lastDiscoveryNodes.getLocalNodeId()).orElse(initialNodeId);
         Task task = request.createTask(
-            lastDiscoveryNodes.getLocalNodeId(),
-            taskIdGenerator.incrementAndGet(),
+            new TaskId(localNodeId, taskIdGenerator.incrementAndGet()),
             type,
             action,
             request.getParentTask(),
diff --git a/server/src/test/java/org/elasticsearch/action/admin/cluster/node/tasks/TaskManagerTestCase.java b/server/src/test/java/org/elasticsearch/action/admin/cluster/node/tasks/TaskManagerTestCase.java
@@ -177,7 +177,7 @@ public TestNode(String name, ThreadPool threadPool, Settings settings) {
                 discoveryNode.set(DiscoveryNodeUtils.create(name, address.publishAddress(), emptyMap(), emptySet()));
                 return discoveryNode.get();
             };
-            TaskManager taskManager = createTaskManager(settings, threadPool, emptySet(), Tracer.NOOP);
+            TaskManager taskManager = createTaskManager(settings, threadPool, emptySet(), Tracer.NOOP, name);
             transportService = new TransportService(
                 settings,
                 new Netty4Transport(
diff --git a/server/src/test/java/org/elasticsearch/action/resync/TransportResyncReplicationActionTests.java b/server/src/test/java/org/elasticsearch/action/resync/TransportResyncReplicationActionTests.java
@@ -123,7 +123,8 @@ public void testResyncDoesNotBlockOnPrimaryAction() throws Exception {
                     NOOP_TRANSPORT_INTERCEPTOR,
                     x -> clusterService.localNode(),
                     null,
-                    Collections.emptySet()
+                    Collections.emptySet(),
+                    clusterService.localNode().getId()
                 );
                 transportService.start();
                 transportService.acceptIncomingRequests();
diff --git a/server/src/test/java/org/elasticsearch/action/support/replication/TransportReplicationActionTests.java b/server/src/test/java/org/elasticsearch/action/support/replication/TransportReplicationActionTests.java
@@ -1316,7 +1316,9 @@ public void testRetryOnReplicaWithRealTransport() throws Exception {
             TransportService.NOOP_TRANSPORT_INTERCEPTOR,
             x -> clusterService.localNode(),
             null,
-            Collections.emptySet()
+            Collections.emptySet(),
+            clusterService.localNode().getId()
+
         );
         transportService.start();
         transportService.acceptIncomingRequests();
diff --git a/server/src/test/java/org/elasticsearch/discovery/FileBasedSeedHostsProviderTests.java b/server/src/test/java/org/elasticsearch/discovery/FileBasedSeedHostsProviderTests.java
@@ -92,6 +92,8 @@ public BoundTransportAddress boundAddress() {
             transport,
             threadPool,
             TransportService.NOOP_TRANSPORT_INTERCEPTOR,
+            null,
+            null,
             null
         );
     }
diff --git a/server/src/test/java/org/elasticsearch/transport/TransportServiceHandshakeTests.java b/server/src/test/java/org/elasticsearch/transport/TransportServiceHandshakeTests.java
@@ -92,7 +92,8 @@ private TransportService startServices(
                 .version(nodeVersion)
                 .build(),
             null,
-            Collections.emptySet()
+            Collections.emptySet(),
+            nodeNameAndId
         );
         transportService.start();
         transportService.acceptIncomingRequests();
diff --git a/test/framework/src/main/java/org/elasticsearch/node/MockNode.java b/test/framework/src/main/java/org/elasticsearch/node/MockNode.java
@@ -167,8 +167,10 @@ protected TransportService newTransportService(
             Function<BoundTransportAddress, DiscoveryNode> localNodeFactory,
             ClusterSettings clusterSettings,
             TaskManager taskManager,
-            Tracer tracer
+            Tracer tracer,
+            String initialNodeId
         ) {
+
             // we use the MockTransportService.TestPlugin class as a marker to create a network
             // module with this MockNetworkService. NetworkService is such an integral part of the systme
             // we don't allow to plug it in from plugins or anything. this is a test-only override and
@@ -183,7 +185,8 @@ protected TransportService newTransportService(
                     localNodeFactory,
                     clusterSettings,
                     taskManager,
-                    tracer
+                    tracer,
+                    initialNodeId
                 );
             } else {
                 return new MockTransportService(
@@ -193,7 +196,8 @@ protected TransportService newTransportService(
                     interceptor,
                     localNodeFactory,
                     clusterSettings,
-                    taskManager.getTaskHeaders()
+                    taskManager.getTaskHeaders(),
+                    initialNodeId
                 );
             }
         }
diff --git a/test/framework/src/main/java/org/elasticsearch/test/ClusterServiceUtils.java b/test/framework/src/main/java/org/elasticsearch/test/ClusterServiceUtils.java
@@ -160,7 +160,7 @@ public static ClusterService createClusterService(
             settings,
             clusterSettings,
             threadPool,
-            new TaskManager(settings, threadPool, Collections.emptySet(), Tracer.NOOP)
+            new TaskManager(settings, threadPool, Collections.emptySet(), Tracer.NOOP, localNode.getId())
         );
         clusterService.setNodeConnectionsService(createNoOpNodeConnectionsService());
         ClusterState.Builder builder = ClusterState.builder(new ClusterName(ClusterServiceUtils.class.getSimpleName()))
diff --git a/test/framework/src/main/java/org/elasticsearch/test/transport/MockTransportService.java b/test/framework/src/main/java/org/elasticsearch/test/transport/MockTransportService.java
@@ -188,20 +188,21 @@ public static MockTransportService createNewService(
         Set<String> taskHeaders,
         TransportInterceptor interceptor
     ) {
+        String initialNodeId = UUIDs.randomBase64UUID();
         return new MockTransportService(
             settings,
             new StubbableTransport(transport),
             threadPool,
             interceptor,
-            boundAddress -> DiscoveryNodeUtils.builder(UUIDs.randomBase64UUID())
+            boundAddress -> DiscoveryNodeUtils.builder(initialNodeId)
                 .name(Node.NODE_NAME_SETTING.get(settings))
                 .address(boundAddress.publishAddress())
                 .attributes(Node.NODE_ATTRIBUTES.getAsMap(settings))
                 .roles(DiscoveryNode.getRolesFromSettings(settings))
                 .version(version)
                 .build(),
             clusterSettings,
-            createTaskManager(settings, threadPool, taskHeaders, Tracer.NOOP)
+            createTaskManager(settings, threadPool, taskHeaders, Tracer.NOOP, initialNodeId)
         );
     }
 
@@ -223,24 +224,25 @@ public static MockTransportService getInstance(String nodeName) {
      *                        updates for {@link TransportSettings#TRACE_LOG_EXCLUDE_SETTING} and
      *                        {@link TransportSettings#TRACE_LOG_INCLUDE_SETTING}.
      */
-    public MockTransportService(
+    public static MockTransportService createMockTransportService(
         Settings settings,
         Transport transport,
         ThreadPool threadPool,
         TransportInterceptor interceptor,
         @Nullable ClusterSettings clusterSettings
     ) {
-        this(
+        String initialNodeId = settings.get(Node.NODE_NAME_SETTING.getKey(), UUIDs.randomBase64UUID());
+        return new MockTransportService(
             settings,
             new StubbableTransport(transport),
             threadPool,
             interceptor,
-            (boundAddress) -> DiscoveryNodeUtils.builder(settings.get(Node.NODE_NAME_SETTING.getKey(), UUIDs.randomBase64UUID()))
+            (boundAddress) -> DiscoveryNodeUtils.builder(initialNodeId)
                 .applySettings(settings)
                 .address(boundAddress.publishAddress())
                 .build(),
             clusterSettings,
-            createTaskManager(settings, threadPool, Set.of(), Tracer.NOOP)
+            createTaskManager(settings, threadPool, Set.of(), Tracer.NOOP, initialNodeId)
         );
     }
 
@@ -258,7 +260,8 @@ public MockTransportService(
         TransportInterceptor interceptor,
         Function<BoundTransportAddress, DiscoveryNode> localNodeFactory,
         @Nullable ClusterSettings clusterSettings,
-        Set<String> taskHeaders
+        Set<String> taskHeaders,
+        @Nullable String initialNodeId
     ) {
         this(
             settings,
@@ -267,7 +270,7 @@ public MockTransportService(
             interceptor,
             localNodeFactory,
             clusterSettings,
-            createTaskManager(settings, threadPool, taskHeaders, Tracer.NOOP)
+            createTaskManager(settings, threadPool, taskHeaders, Tracer.NOOP, initialNodeId)
         );
     }
 
@@ -277,7 +280,8 @@ public MockTransportService(
         ThreadPool threadPool,
         TransportInterceptor interceptor,
         Function<BoundTransportAddress, DiscoveryNode> localNodeFactory,
-        @Nullable ClusterSettings clusterSettings
+        @Nullable ClusterSettings clusterSettings,
+        String initialNodeId
     ) {
         this(
             settings,
@@ -286,7 +290,7 @@ public MockTransportService(
             interceptor,
             localNodeFactory,
             clusterSettings,
-            createTaskManager(settings, threadPool, Set.of(), Tracer.NOOP)
+            createTaskManager(settings, threadPool, Set.of(), Tracer.NOOP, initialNodeId)
         );
     }
 
@@ -331,19 +335,28 @@ private static TransportAddress[] extractTransportAddresses(TransportService tra
         return transportAddresses.toArray(new TransportAddress[transportAddresses.size()]);
     }
 
-    public static TaskManager createTaskManager(Settings settings, ThreadPool threadPool, Set<String> taskHeaders, Tracer tracer) {
-        if (MockTaskManager.SPY_TASK_MANAGER_SETTING.get(settings)) {
-            return spy(createMockTaskManager(settings, threadPool, taskHeaders, tracer));
-        } else {
-            return createMockTaskManager(settings, threadPool, taskHeaders, tracer);
-        }
+    public static TaskManager createTaskManager(
+        Settings settings,
+        ThreadPool threadPool,
+        Set<String> taskHeaders,
+        Tracer tracer,
+        @Nullable String initialNodeId
+    ) {
+        TaskManager mockTaskManager = createMockTaskManager(settings, threadPool, taskHeaders, tracer, initialNodeId);
+        return MockTaskManager.SPY_TASK_MANAGER_SETTING.get(settings) ? spy(mockTaskManager) : mockTaskManager;
     }
 
-    private static TaskManager createMockTaskManager(Settings settings, ThreadPool threadPool, Set<String> taskHeaders, Tracer tracer) {
+    private static TaskManager createMockTaskManager(
+        Settings settings,
+        ThreadPool threadPool,
+        Set<String> taskHeaders,
+        Tracer tracer,
+        @Nullable String initialNodeId
+    ) {
         if (MockTaskManager.USE_MOCK_TASK_MANAGER_SETTING.get(settings)) {
             return new MockTaskManager(settings, threadPool, taskHeaders);
         } else {
-            return new TaskManager(settings, threadPool, taskHeaders, tracer);
+            return new TaskManager(settings, threadPool, taskHeaders, tracer, initialNodeId);
         }
     }
 
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryRequest.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryRequest.java
@@ -248,9 +248,9 @@ public EsqlQueryRequest allowPartialResults(boolean allowPartialResults) {
     }
 
     @Override
-    public Task createTask(String localNodeId, long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
-        var status = new EsqlQueryStatus(new AsyncExecutionId(UUIDs.randomBase64UUID(), new TaskId(localNodeId, id)));
-        return new EsqlQueryRequestTask(query, id, type, action, parentTaskId, headers, status);
+    public Task createTask(TaskId taskId, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+        var status = new EsqlQueryStatus(new AsyncExecutionId(UUIDs.randomBase64UUID(), taskId));
+        return new EsqlQueryRequestTask(query, taskId.getId(), type, action, parentTaskId, headers, status);
     }
 
     private static class EsqlQueryRequestTask extends CancellableTask {
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/action/EsqlQueryRequestTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/action/EsqlQueryRequestTests.java
@@ -611,7 +611,7 @@ public void testTask() throws IOException {
 
         EsqlQueryRequest request = parseEsqlQueryRequestSync(requestJson);
         String localNode = randomAlphaOfLength(2);
-        Task task = request.createTask(localNode, id, "transport", EsqlQueryAction.NAME, TaskId.EMPTY_TASK_ID, Map.of());
+        Task task = request.createTask(new TaskId(localNode, id), "transport", EsqlQueryAction.NAME, TaskId.EMPTY_TASK_ID, Map.of());
         assertThat(task.getDescription(), equalTo(query));
 
         TaskInfo taskInfo = task.taskInfo(localNode, true);

Original file line number	Diff line number	Diff line change
`@@ -56,8 +56,8 @@ public boolean shouldEnable() {`
`56`	`56`	`}`
`57`	`57`
`58`	`58`	`@Override`
`59`		`- public Task createTask(String localNodeId, long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {`
`60`		`- return new CancellableTask(id, type, action, "Logs streams activation toggle request", parentTaskId, headers);`
	`59`	`+ public Task createTask(TaskId taskId, String type, String action, TaskId parentTaskId, Map<String, String> headers) {`
	`60`	`+ return new CancellableTask(taskId.getId(), type, action, "Logs streams activation toggle request", parentTaskId, headers);`
`61`	`61`	`}`
`62`	`62`	`}`
`63`	`63`	`}`
Original file line number	Diff line number	Diff line change
`@@ -61,6 +61,8 @@ protected MockTransportService createTransportService() {`
`61`	`61`	`),`
`62`	`62`	`threadPool,`
`63`	`63`	`TransportService.NOOP_TRANSPORT_INTERCEPTOR,`
	`64`	`+ null,`
	`65`	`+ null,`
`64`	`66`	`null`
`65`	`67`	`);`
`66`	`68`	`}`