From bd8470a9bcb6e2f74a620a72580a9df67f91b2ee Mon Sep 17 00:00:00 2001
From: Dejan Mircevski <dejan.mircevski@starburstdata.com>
Date: Tue, 16 Jan 2024 21:42:37 -0500
Subject: [PATCH] Pass InternalNode to HttpRemoteTask constructor

This allows us to extract more info than just the node ID in the future.
---
 .../main/java/io/trino/server/HttpRemoteTaskFactory.java   | 2 +-
 .../java/io/trino/server/remotetask/HttpRemoteTask.java    | 7 ++++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/core/trino-main/src/main/java/io/trino/server/HttpRemoteTaskFactory.java b/core/trino-main/src/main/java/io/trino/server/HttpRemoteTaskFactory.java
index 31c7916576b0..936ff86bcd0b 100644
--- a/core/trino-main/src/main/java/io/trino/server/HttpRemoteTaskFactory.java
+++ b/core/trino-main/src/main/java/io/trino/server/HttpRemoteTaskFactory.java
@@ -153,7 +153,7 @@ public RemoteTask createRemoteTask(
                 session,
                 stageSpan,
                 taskId,
-                node.getNodeIdentifier(),
+                node,
                 speculative,
                 locationFactory.createTaskLocation(node, taskId),
                 fragment,
diff --git a/core/trino-main/src/main/java/io/trino/server/remotetask/HttpRemoteTask.java b/core/trino-main/src/main/java/io/trino/server/remotetask/HttpRemoteTask.java
index 5bd5ddc80a65..faf0c34e25d3 100644
--- a/core/trino-main/src/main/java/io/trino/server/remotetask/HttpRemoteTask.java
+++ b/core/trino-main/src/main/java/io/trino/server/remotetask/HttpRemoteTask.java
@@ -57,6 +57,7 @@
 import io.trino.execution.buffer.PipelinedBufferInfo;
 import io.trino.execution.buffer.PipelinedOutputBuffers;
 import io.trino.execution.buffer.SpoolingOutputStats;
+import io.trino.metadata.InternalNode;
 import io.trino.metadata.Split;
 import io.trino.operator.RetryPolicy;
 import io.trino.operator.TaskStats;
@@ -210,7 +211,7 @@ public HttpRemoteTask(
             Session session,
             Span stageSpan,
             TaskId taskId,
-            String nodeId,
+            InternalNode node,
             boolean speculative,
             URI location,
             PlanFragment planFragment,
@@ -240,7 +241,7 @@ public HttpRemoteTask(
         requireNonNull(session, "session is null");
         requireNonNull(stageSpan, "stageSpan is null");
         requireNonNull(taskId, "taskId is null");
-        requireNonNull(nodeId, "nodeId is null");
+        requireNonNull(node, "node is null");
         requireNonNull(location, "location is null");
         requireNonNull(planFragment, "planFragment is null");
         requireNonNull(outputBuffers, "outputBuffers is null");
@@ -258,7 +259,7 @@ public HttpRemoteTask(
             this.taskId = taskId;
             this.session = session;
             this.stageSpan = stageSpan;
-            this.nodeId = nodeId;
+            this.nodeId = node.getNodeIdentifier();
             this.speculative = new AtomicBoolean(speculative);
             this.planFragment = planFragment;
             this.outputBuffers.set(outputBuffers);