Teach tasks what machine they are running on

core/src/main/java/org/apache/druid/guice/annotations/Parent.java

@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.guice.annotations;
+
+import com.google.inject.BindingAnnotation;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Binding annotation for org.apache.druid.server.DruidNode.
+ * Indicates that the DruidNode bound with this annotation holds the information of the machine which spawned
+ * this process.
+ *
+ * @see Self
+ */
+@Target({ElementType.FIELD, ElementType.PARAMETER})
+@Retention(RetentionPolicy.RUNTIME)
+@BindingAnnotation
+@PublicApi
+public @interface Parent
+{
+}

core/src/main/java/org/apache/druid/guice/annotations/Self.java

@@ -27,6 +27,11 @@
 import java.lang.annotation.Target;
 
 /**
+ * Binding annotation for org.apache.druid.server.DruidNode.
+ * Indicates that the DruidNode bound with this annotation holds the information of the machine where this process
+ * is running.
+ *
+ * @see Parent
  */
 @Target({ElementType.FIELD, ElementType.PARAMETER, ElementType.METHOD})
 @Retention(RetentionPolicy.RUNTIME)

extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java

@@ -2709,6 +2709,7 @@ public void close()
 
     toolboxFactory = new TaskToolboxFactory(
         taskConfig,
+        null, // taskExecutorNode
         taskActionClientFactory,
         emitter,
         dataSegmentPusher,

extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskTest.java

@@ -2889,6 +2889,7 @@ public void close()
 
     toolboxFactory = new TaskToolboxFactory(
         taskConfig,
+        null, // taskExecutorNode
         taskActionClientFactory,
         emitter,
         dataSegmentPusher,

indexing-service/src/main/java/org/apache/druid/indexing/common/TaskToolbox.java

@@ -67,6 +67,7 @@
 public class TaskToolbox
 {
   private final TaskConfig config;
+  private final DruidNode taskExecutorNode;
   private final TaskActionClient taskActionClient;
   private final ServiceEmitter emitter;
   private final DataSegmentPusher segmentPusher;
@@ -101,6 +102,7 @@ public class TaskToolbox
 
   public TaskToolbox(
       TaskConfig config,
+      DruidNode taskExecutorNode,
       TaskActionClient taskActionClient,
       ServiceEmitter emitter,
       DataSegmentPusher segmentPusher,
@@ -129,6 +131,7 @@ public TaskToolbox(
   )
   {
     this.config = config;
+    this.taskExecutorNode = taskExecutorNode;
     this.taskActionClient = taskActionClient;
     this.emitter = emitter;
     this.segmentPusher = segmentPusher;
@@ -162,6 +165,11 @@ public TaskConfig getConfig()
     return config;
   }
 
+  public DruidNode getTaskExecutorNode()
+  {
+    return taskExecutorNode;
+  }
+
   public TaskActionClient getTaskActionClient()
   {
     return taskActionClient;

indexing-service/src/main/java/org/apache/druid/indexing/common/TaskToolboxFactory.java

@@ -29,6 +29,7 @@
 import org.apache.druid.discovery.DataNodeService;
 import org.apache.druid.discovery.DruidNodeAnnouncer;
 import org.apache.druid.discovery.LookupNodeService;
+import org.apache.druid.guice.annotations.Parent;
 import org.apache.druid.guice.annotations.Processing;
 import org.apache.druid.guice.annotations.RemoteChatHandler;
 import org.apache.druid.indexing.common.actions.TaskActionClientFactory;
@@ -57,6 +58,7 @@
 public class TaskToolboxFactory
 {
   private final TaskConfig config;
+  private final DruidNode taskExecutorNode;
   private final TaskActionClientFactory taskActionClientFactory;
   private final ServiceEmitter emitter;
   private final DataSegmentPusher segmentPusher;
@@ -85,6 +87,7 @@ public class TaskToolboxFactory
   @Inject
   public TaskToolboxFactory(
       TaskConfig config,
+      @Parent DruidNode taskExecutorNode,
       TaskActionClientFactory taskActionClientFactory,
       ServiceEmitter emitter,
       DataSegmentPusher segmentPusher,
@@ -112,6 +115,7 @@ public TaskToolboxFactory(
   )
   {
     this.config = config;
+    this.taskExecutorNode = taskExecutorNode;
     this.taskActionClientFactory = taskActionClientFactory;
     this.emitter = emitter;
     this.segmentPusher = segmentPusher;
@@ -143,6 +147,7 @@ public TaskToolbox build(Task task)
     final File taskWorkDir = config.getTaskWorkDir(task.getId());
     return new TaskToolbox(
         config,
+        taskExecutorNode,
         taskActionClientFactory.create(task),
         emitter,
         segmentPusher,

indexing-service/src/main/java/org/apache/druid/indexing/overlord/ForkingTaskRunner.java

@@ -289,14 +289,31 @@ public TaskStatus call()
                         command.add(StringUtils.format("-Ddruid.host=%s", childHost));
                         command.add(StringUtils.format("-Ddruid.plaintextPort=%d", childPort));
                         command.add(StringUtils.format("-Ddruid.tlsPort=%d", tlsChildPort));
-                        /**
-                         * These are not enabled per default to allow the user to either set or not set them
-                         * Users are highly suggested to be set in druid.indexer.runner.javaOpts
-                         * See org.apache.druid.concurrent.TaskThreadPriority#getThreadPriorityFromTaskPriority(int)
-                         * for more information
-                         command.add("-XX:+UseThreadPriorities");
-                         command.add("-XX:ThreadPriorityPolicy=42");
-                         */
+
+                        // Let tasks know where they are running on.
+                        // This information is used in native parallel indexing with shuffle.
+                        command.add(StringUtils.format("-Ddruid.task.executor.service=%s", node.getServiceName()));
+                        command.add(StringUtils.format("-Ddruid.task.executor.host=%s", node.getHost()));
+                        command.add(
+                            StringUtils.format("-Ddruid.task.executor.plaintextPort=%d", node.getPlaintextPort())
+                        );
+                        command.add(
+                            StringUtils.format(
+                                "-Ddruid.task.executor.enablePlaintextPort=%s",
+                                node.isEnablePlaintextPort()
+                            )
+                        );
+                        command.add(StringUtils.format("-Ddruid.task.executor.tlsPort=%d", node.getTlsPort()));
+                        command.add(
+                            StringUtils.format("-Ddruid.task.executor.enableTlsPort=%s", node.isEnableTlsPort())
+                        );
+
+                        // These are not enabled per default to allow the user to either set or not set them
+                        // Users are highly suggested to be set in druid.indexer.runner.javaOpts
+                        // See org.apache.druid.concurrent.TaskThreadPriority#getThreadPriorityFromTaskPriority(int)
+                        // for more information
+                        // command.add("-XX:+UseThreadPriorities");
+                        // command.add("-XX:ThreadPriorityPolicy=42");
 
                         command.add("org.apache.druid.cli.Main");
                         command.add("internal");

indexing-service/src/test/java/org/apache/druid/indexing/common/TaskToolboxTest.java

@@ -41,6 +41,7 @@
 import org.apache.druid.segment.loading.SegmentLoaderLocalCacheManager;
 import org.apache.druid.segment.loading.SegmentLoadingException;
 import org.apache.druid.segment.realtime.plumber.SegmentHandoffNotifierFactory;
+import org.apache.druid.server.DruidNode;
 import org.apache.druid.server.coordination.DataSegmentAnnouncer;
 import org.apache.druid.server.coordination.DataSegmentServerAnnouncer;
 import org.apache.druid.timeline.DataSegment;
@@ -96,6 +97,7 @@ public void setUp() throws IOException
 
     taskToolbox = new TaskToolboxFactory(
         new TaskConfig(temporaryFolder.newFile().toString(), null, null, 50000, null, false, null, null, null),
+        new DruidNode("druid/middlemanager", "localhost", false, 8091, null, true, false),
         mockTaskActionClientFactory,
         mockEmitter,
         mockSegmentPusher,

indexing-service/src/test/java/org/apache/druid/indexing/common/task/AppenderatorDriverRealtimeIndexTaskTest.java

@@ -1602,6 +1602,7 @@ public List<StorageLocationConfig> getLocations()
 
     taskToolboxFactory = new TaskToolboxFactory(
         taskConfig,
+        new DruidNode("druid/middlemanager", "localhost", false, 8091, null, true, false),
         taskActionClientFactory,
         emitter,
         new TestDataSegmentPusher(),

indexing-service/src/test/java/org/apache/druid/indexing/common/task/CompactionTaskRunTest.java

@@ -696,6 +696,7 @@ public List<StorageLocationConfig> getLocations()
     );
 
     final TaskToolbox box = new TaskToolbox(
+        null,
         null,
         actionClient,
         null,

indexing-service/src/test/java/org/apache/druid/indexing/common/task/CompactionTaskTest.java

@@ -1160,6 +1160,7 @@ private static class TestTaskToolbox extends TaskToolbox
     )
     {
       super(
+          null,
           null,
           taskActionClient,
           null,

indexing-service/src/test/java/org/apache/druid/indexing/common/task/IngestionTestBase.java

@@ -58,6 +58,7 @@
 import org.apache.druid.segment.loading.LocalDataSegmentPusher;
 import org.apache.druid.segment.loading.LocalDataSegmentPusherConfig;
 import org.apache.druid.segment.loading.NoopDataSegmentKiller;
+import org.apache.druid.server.DruidNode;
 import org.apache.druid.server.metrics.NoopServiceEmitter;
 import org.apache.druid.timeline.DataSegment;
 import org.junit.After;
@@ -278,6 +279,7 @@ public ListenableFuture<TaskStatus> run(Task task)
 
         final TaskToolbox box = new TaskToolbox(
             null,
+            new DruidNode("druid/middlemanager", "localhost", false, 8091, null, true, false),
             taskActionClient,
             null,
             new LocalDataSegmentPusher(new LocalDataSegmentPusherConfig()),

indexing-service/src/test/java/org/apache/druid/indexing/common/task/RealtimeIndexTaskTest.java

@@ -976,6 +976,7 @@ public List<StorageLocationConfig> getLocations()
     };
     final TaskToolboxFactory toolboxFactory = new TaskToolboxFactory(
         taskConfig,
+        null, // taskExecutorNode
         taskActionClientFactory,
         emitter,
         new TestDataSegmentPusher(),

indexing-service/src/test/java/org/apache/druid/indexing/common/task/batch/parallel/AbstractParallelIndexSupervisorTaskTest.java

@@ -55,6 +55,7 @@
 import org.apache.druid.segment.loading.NoopDataSegmentKiller;
 import org.apache.druid.segment.realtime.appenderator.SegmentIdWithShardSpec;
 import org.apache.druid.segment.realtime.firehose.NoopChatHandlerProvider;
+import org.apache.druid.server.DruidNode;
 import org.apache.druid.server.security.AllowAllAuthorizer;
 import org.apache.druid.server.security.Authorizer;
 import org.apache.druid.server.security.AuthorizerMapper;
@@ -219,6 +220,7 @@ TaskToolbox createTaskToolbox(Task task) throws IOException
   {
     return new TaskToolbox(
         null,
+        new DruidNode("druid/middlemanager", "localhost", false, 8091, null, true, false),
         actionClient,
         null,
         new LocalDataSegmentPusher(

indexing-service/src/test/java/org/apache/druid/indexing/overlord/SingleTaskBackgroundRunnerTest.java

@@ -82,6 +82,7 @@ public void setup() throws IOException
     final ServiceEmitter emitter = new NoopServiceEmitter();
     final TaskToolboxFactory toolboxFactory = new TaskToolboxFactory(
         taskConfig,
+        null,
         EasyMock.createMock(TaskActionClientFactory.class),
         emitter,
         new NoopDataSegmentPusher(),

indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskLifecycleTest.java

@@ -560,6 +560,7 @@ public List<StorageLocationConfig> getLocations()
 
     return new TaskToolboxFactory(
         taskConfig,
+        new DruidNode("druid/middlemanager", "localhost", false, 8091, null, true, false),
         tac,
         emitter,
         dataSegmentPusher,

indexing-service/src/test/java/org/apache/druid/indexing/worker/WorkerTaskManagerTest.java

@@ -108,6 +108,7 @@ public List<StorageLocationConfig> getLocations()
         new TestTaskRunner(
             new TaskToolboxFactory(
                 taskConfig,
+                null,
                 taskActionClientFactory,
                 null,
                 null,

indexing-service/src/test/java/org/apache/druid/indexing/worker/WorkerTaskMonitorTest.java

@@ -166,6 +166,7 @@ private WorkerTaskMonitor createTaskMonitor()
         new SingleTaskBackgroundRunner(
             new TaskToolboxFactory(
                 taskConfig,
+                null,
                 taskActionClientFactory,
                 null, null, null, null, null, null, null, notifierFactory, null, null, null,
                 new SegmentLoaderFactory(null, jsonMapper),

services/src/main/java/org/apache/druid/cli/CliIndexer.java

@@ -41,6 +41,7 @@
 import org.apache.druid.guice.QueryRunnerFactoryModule;
 import org.apache.druid.guice.QueryableModule;
 import org.apache.druid.guice.QueryablePeonModule;
+import org.apache.druid.guice.annotations.Parent;
 import org.apache.druid.guice.annotations.RemoteChatHandler;
 import org.apache.druid.guice.annotations.Self;
 import org.apache.druid.indexing.common.MultipleFileTaskReportFileWriter;
@@ -103,6 +104,7 @@ public void configure(Binder binder)
 
             IndexingServiceModuleHelper.configureTaskRunnerConfigs(binder);
 
+            JsonConfigProvider.bind(binder, "druid", DruidNode.class, Parent.class);
             JsonConfigProvider.bind(binder, "druid.worker", WorkerConfig.class);
 
             CliPeon.bindTaskConfigAndClients(binder);

services/src/main/java/org/apache/druid/cli/CliPeon.java

@@ -55,6 +55,7 @@
 import org.apache.druid.guice.QueryableModule;
 import org.apache.druid.guice.QueryablePeonModule;
 import org.apache.druid.guice.annotations.Json;
+import org.apache.druid.guice.annotations.Parent;
 import org.apache.druid.guice.annotations.Smile;
 import org.apache.druid.indexing.common.RetryPolicyConfig;
 import org.apache.druid.indexing.common.RetryPolicyFactory;
@@ -100,6 +101,7 @@
 import org.apache.druid.segment.realtime.plumber.CoordinatorBasedSegmentHandoffNotifierConfig;
 import org.apache.druid.segment.realtime.plumber.CoordinatorBasedSegmentHandoffNotifierFactory;
 import org.apache.druid.segment.realtime.plumber.SegmentHandoffNotifierFactory;
+import org.apache.druid.server.DruidNode;
 import org.apache.druid.server.coordination.BatchDataSegmentAnnouncer;
 import org.apache.druid.server.coordination.ServerType;
 import org.apache.druid.server.http.SegmentListerResource;
@@ -169,6 +171,8 @@ public void configure(Binder binder)
             binder.bindConstant().annotatedWith(Names.named("servicePort")).to(0);
             binder.bindConstant().annotatedWith(Names.named("tlsServicePort")).to(-1);
 
+            JsonConfigProvider.bind(binder, "druid.task.executor", DruidNode.class, Parent.class);
+
             bindRowIngestionMeters(binder);
 
             bindChatHandler(binder);