feat: Multi-cluster Kubernetes Task Runner

extensions-core/kubernetes-overlord-extensions/pom.xml

@@ -77,6 +77,10 @@
       <artifactId>commons-lang3</artifactId>
       <scope>provided</scope>
     </dependency>
+    <dependency>
+      <groupId>com.fasterxml.jackson.dataformat</groupId>
+      <artifactId>jackson-dataformat-properties</artifactId>
+    </dependency>
     <dependency>
       <groupId>io.netty</groupId>
       <artifactId>netty</artifactId>

extensions-core/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/AutoscalableThreadPoolExecutor.java

@@ -0,0 +1,122 @@
+/*
+ * 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.k8s.overlord;
+
+
+import org.apache.druid.common.config.ConfigManager;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.common.concurrent.Execs;
+import org.apache.druid.java.util.emitter.EmittingLogger;
+import org.apache.druid.k8s.overlord.execution.KubernetesTaskRunnerDynamicConfig;
+
+import javax.annotation.Nullable;
+import java.util.List;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
+public class AutoscalableThreadPoolExecutor extends ThreadPoolExecutor
+{
+  private static final EmittingLogger log = new EmittingLogger(AutoscalableThreadPoolExecutor.class);
+
+  @Nullable
+  private final ConfigManager configManager;
+  private final String listenerKey;
+  private final Consumer<KubernetesTaskRunnerDynamicConfig> configListener;
+
+  public AutoscalableThreadPoolExecutor(int initialCapacity, @Nullable ConfigManager configManager)
+  {
+    super(
+        initialCapacity,
+        initialCapacity,
+        0L,
+        TimeUnit.MILLISECONDS,
+        new LinkedBlockingQueue<>(),
+        Execs.makeThreadFactory("k8s-task-runner-%d", null)
+    );
+
+    this.configManager = configManager;
+    this.listenerKey = StringUtils.format("AutoscalableThreadPoolExecutor-%d", System.identityHashCode(this));
+    this.configListener = this::onConfigurationChange;
+
+    // Monitor the configuration change
+    if (configManager != null && !configManager.addListener(
+        KubernetesTaskRunnerDynamicConfig.CONFIG_KEY,
+        listenerKey,
+        configListener)) {
+      log.error("Failed to add configuration listener for AutoscalableThreadPoolExecutor");
+    }
+  }
+
+  @Override
+  public void shutdown()
+  {
+    removeConfigListener();
+    super.shutdown();
+  }
+
+  @Override
+  public List<Runnable> shutdownNow()
+  {
+    removeConfigListener();
+    return super.shutdownNow();
+  }
+
+  private void removeConfigListener()
+  {
+    if (configManager == null) {
+      return;
+    }
+
+    if (!configManager.removeListener(
+        KubernetesTaskRunnerDynamicConfig.CONFIG_KEY,
+        listenerKey,
+        configListener
+    )) {
+      log.warn("Failed to remove configuration listener[%s] for AutoscalableThreadPoolExecutor", listenerKey);
+    }
+  }
+
+  private void onConfigurationChange(KubernetesTaskRunnerDynamicConfig config)
+  {
+    if (config.getCapacity() == null) {
+      return;
+    }
+
+    final int curCapacity = this.getCorePoolSize();
+    final int newCapacity = config.getCapacity();
+    if (newCapacity == curCapacity) {
+      return;
+    }
+
+    log.info("Adjusting k8s task runner capacity from [%d] to [%d]", curCapacity, newCapacity);
+    // maximum pool size must always be greater than or equal to the core pool size
+    if (newCapacity < curCapacity) {
+      // decrease capacity
+      setCorePoolSize(newCapacity);
+      setMaximumPoolSize(newCapacity);
+    } else {
+      // increase capacity
+      setMaximumPoolSize(newCapacity);
+      setCorePoolSize(newCapacity);
+    }
+  }
+}

extensions-core/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesOverlordModule.java

@@ -119,8 +119,12 @@ public void configure(Binder binder)
     biddy.addBinding(KubernetesAndWorkerTaskRunnerFactory.TYPE_NAME)
          .to(KubernetesAndWorkerTaskRunnerFactory.class)
          .in(LazySingleton.class);
+    biddy.addBinding(MultipleKubernetesTaskRunnerFactory.TYPE_NAME)
+         .to(MultipleKubernetesTaskRunnerFactory.class)
+         .in(LazySingleton.class);
     binder.bind(KubernetesTaskRunnerFactory.class).in(LazySingleton.class);
     binder.bind(KubernetesAndWorkerTaskRunnerFactory.class).in(LazySingleton.class);
+    binder.bind(MultipleKubernetesTaskRunnerFactory.class).in(LazySingleton.class);
     binder.bind(RunnerStrategy.class)
           .toProvider(RunnerStrategyProvider.class)
           .in(LazySingleton.class);

extensions-core/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java

@@ -122,6 +122,7 @@
   private final HttpClient httpClient;
   private final PeonLifecycleFactory peonLifecycleFactory;
   private final ServiceEmitter emitter;
+  private final boolean ownsExecutor;
   // currently worker categories aren't supported, so it's hardcoded.
   protected static final String WORKER_CATEGORY = "_k8s_worker_category";
 
@@ -136,6 +137,20 @@
       ServiceEmitter emitter,
       ConfigManager configManager
   )
+  {
+    this(adapter, config, client, httpClient, peonLifecycleFactory, emitter, null, configManager);
+  }
+
+  public KubernetesTaskRunner(
+      TaskAdapter adapter,
+      KubernetesTaskRunnerConfig config,
+      KubernetesPeonClient client,
+      HttpClient httpClient,
+      PeonLifecycleFactory peonLifecycleFactory,
+      ServiceEmitter emitter,
+      @Nullable ThreadPoolExecutor sharedExecutor,
+      @Nullable ConfigManager configManager
+  )
   {
     this.adapter = adapter;
     this.config = config;
@@ -146,9 +161,28 @@
     this.emitter = emitter;
 
     this.currentCapacity = new AtomicInteger(config.getCapacity());
-    this.tpe = new ThreadPoolExecutor(currentCapacity.get(), currentCapacity.get(), 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<Runnable>(), Execs.makeThreadFactory("k8s-task-runner-%d", null));
+    if (sharedExecutor == null) {
+      this.tpe = new ThreadPoolExecutor(
+          currentCapacity.get(),
+          currentCapacity.get(),
+          0L,
+          TimeUnit.MILLISECONDS,
+          new LinkedBlockingQueue<Runnable>(),
+          Execs.makeThreadFactory("k8s-task-runner-%d", null)
+      );
+      this.ownsExecutor = true;
+    } else {
+      this.tpe = sharedExecutor;
+      this.ownsExecutor = false;
+    }
     this.exec = MoreExecutors.listeningDecorator(this.tpe);
-    configManager.addListener(KubernetesTaskRunnerDynamicConfig.CONFIG_KEY, StringUtils.format(OBSERVER_KEY, Thread.currentThread().getId()), this::syncCapacityWithDynamicConfig);
+    if (ownsExecutor && configManager != null) {
+      configManager.addListener(
+          KubernetesTaskRunnerDynamicConfig.CONFIG_KEY,
+          StringUtils.format(OBSERVER_KEY, Thread.currentThread().getId()),
+          this::syncCapacityWithDynamicConfig
+      );
+    }
   }
 
   @Override
@@ -194,7 +228,7 @@
 
   private void syncCapacityWithDynamicConfig(KubernetesTaskRunnerDynamicConfig config)
   {
-    int newCapacity = config.getCapacity();
+    final int newCapacity = config.getCapacity();
     if (newCapacity == currentCapacity.get()) {
       return;
     }
@@ -210,6 +244,12 @@
     currentCapacity.set(newCapacity);
   }
 
+  @VisibleForTesting
+  KubernetesPeonClient getPeonClient()
+  {
+    return client;
+  }
+
   private TaskStatus runTask(Task task)
   {
     return doTask(task, true);
@@ -443,7 +483,9 @@
   {
     log.info("Stopping KubernetesTaskRunner");
     // Stop managing the running k8s jobs
-    exec.shutdownNow();
+    if (ownsExecutor) {
+      exec.shutdownNow();
+    }
     cleanupExecutor.shutdownNow();
     log.debug("Stopped KubernetesTaskRunner");
   }