Add annotation for pod template (#16772)

* Add annotation for pod template

* pr comments

* add test cases

* add tests

Author: George Shiqi Wu

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/common/DruidK8sConstants.java

@@ -28,6 +28,7 @@ public class DruidK8sConstants
   public static final String TASK_TYPE = "task.type";
   public static final String TASK_GROUP_ID = "task.group.id";
   public static final String TASK_DATASOURCE = "task.datasource";
+  public static final String TASK_JOB_TEMPLATE = "task.jobTemplate";
   public static final int PORT = 8100;
   public static final int TLS_PORT = 8091;
   public static final int DEFAULT_CPU_MILLICORES = 1000;
@@ -42,6 +43,7 @@ public class DruidK8sConstants
   public static final String DRUID_HOSTNAME_ENV = "HOSTNAME";
   public static final String LABEL_KEY = "druid.k8s.peons";
   public static final String DRUID_LABEL_PREFIX = "druid.";
+  public static final String BASE_TEMPLATE_NAME = "base";
   public static final long MAX_ENV_VARIABLE_KBS = 130048; // 127 KB
   static final Predicate<Throwable> IS_TRANSIENT = e -> e instanceof KubernetesResourceNotFoundException;
 }

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/execution/PodTemplateSelectStrategy.java

@@ -23,6 +23,7 @@
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
 import io.fabric8.kubernetes.api.model.PodTemplate;
 import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.k8s.overlord.taskadapter.PodTemplateWithName;
 
 import javax.validation.constraints.NotNull;
 import java.util.Map;
@@ -42,7 +43,7 @@ public interface PodTemplateSelectStrategy
    * allows for customized resource allocation and management tailored to the task's specific requirements.
    *
    * @param task The task for which the Pod template is determined.
-   * @return The pod template that should be used to run the task.
+   * @return The PodTemplateWithName POJO that contains the name of the template selected and the template itself.
    */
-  @NotNull PodTemplate getPodTemplateForTask(Task task, Map<String, PodTemplate> templates);
+  @NotNull PodTemplateWithName getPodTemplateForTask(Task task, Map<String, PodTemplate> templates);
 }

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/execution/SelectorBasedPodTemplateSelectStrategy.java

@@ -24,6 +24,8 @@
 import com.google.common.base.Preconditions;
 import io.fabric8.kubernetes.api.model.PodTemplate;
 import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.k8s.overlord.common.DruidK8sConstants;
+import org.apache.druid.k8s.overlord.taskadapter.PodTemplateWithName;
 
 import java.util.List;
 import java.util.Map;
@@ -53,15 +55,18 @@ public SelectorBasedPodTemplateSelectStrategy(
    * @return the template if a selector matches, otherwise fallback to base template
    */
   @Override
-  public PodTemplate getPodTemplateForTask(Task task, Map<String, PodTemplate> templates)
+  public PodTemplateWithName getPodTemplateForTask(Task task, Map<String, PodTemplate> templates)
   {
     String templateKey = selectors.stream()
                                   .filter(selector -> selector.evaluate(task))
                                   .findFirst()
                                   .map(Selector::getSelectionKey)
-                                  .orElse("base");
+                                  .orElse(DruidK8sConstants.BASE_TEMPLATE_NAME);
 
-    return templates.getOrDefault(templateKey, templates.get("base"));
+    if (!templates.containsKey(templateKey)) {
+      templateKey = DruidK8sConstants.BASE_TEMPLATE_NAME;
+    }
+    return new PodTemplateWithName(templateKey, templates.get(templateKey));
   }
 
   @JsonProperty

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/execution/TaskTypePodTemplateSelectStrategy.java

@@ -22,6 +22,8 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import io.fabric8.kubernetes.api.model.PodTemplate;
 import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.k8s.overlord.common.DruidK8sConstants;
+import org.apache.druid.k8s.overlord.taskadapter.PodTemplateWithName;
 
 import java.util.Map;
 
@@ -40,9 +42,10 @@ public TaskTypePodTemplateSelectStrategy()
   }
 
   @Override
-  public PodTemplate getPodTemplateForTask(Task task, Map<String, PodTemplate> templates)
+  public PodTemplateWithName getPodTemplateForTask(Task task, Map<String, PodTemplate> templates)
   {
-    return templates.getOrDefault(task.getType(), templates.get("base"));
+    String templateKey = templates.containsKey(task.getType()) ? task.getType() : DruidK8sConstants.BASE_TEMPLATE_NAME;
+    return new PodTemplateWithName(templateKey, templates.get(templateKey));
   }
 
   @Override

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/taskadapter/PodTemplateTaskAdapter.java

@@ -138,19 +138,21 @@ public Job fromTask(Task task) throws IOException
       podTemplateSelectStrategy = dynamicConfig.getPodTemplateSelectStrategy();
     }
 
-    PodTemplate podTemplate = podTemplateSelectStrategy.getPodTemplateForTask(task, templates);
+    PodTemplateWithName podTemplateWithName = podTemplateSelectStrategy.getPodTemplateForTask(task, templates);
 
     return new JobBuilder()
         .withNewMetadata()
         .withName(new K8sTaskId(task).getK8sJobName())
         .addToLabels(getJobLabels(taskRunnerConfig, task))
         .addToAnnotations(getJobAnnotations(taskRunnerConfig, task))
+        .addToAnnotations(DruidK8sConstants.TASK_JOB_TEMPLATE, podTemplateWithName.getName())
         .endMetadata()
         .withNewSpec()
-        .withTemplate(podTemplate.getTemplate())
+        .withTemplate(podTemplateWithName.getPodTemplate().getTemplate())
         .editTemplate()
         .editOrNewMetadata()
         .addToAnnotations(getPodTemplateAnnotations(task))
+        .addToAnnotations(DruidK8sConstants.TASK_JOB_TEMPLATE, podTemplateWithName.getName())
         .addToLabels(getPodLabels(taskRunnerConfig, task))
         .endMetadata()
         .editSpec()

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/taskadapter/PodTemplateWithName.java

@@ -0,0 +1,78 @@
+/*
+ * 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.taskadapter;
+
+import io.fabric8.kubernetes.api.model.PodTemplate;
+
+import javax.annotation.Nonnull;
+import java.util.Objects;
+
+public class PodTemplateWithName
+{
+  private final String name;
+  private final PodTemplate podTemplate;
+
+  public PodTemplateWithName(String name, PodTemplate podTemplate)
+  {
+    this.name = name;
+    this.podTemplate = podTemplate;
+  }
+
+  @Nonnull
+  public String getName()
+  {
+    return name;
+  }
+
+  @Nonnull
+  public PodTemplate getPodTemplate()
+  {
+    return podTemplate;
+  }
+
+  @Override
+  public boolean equals(Object o)
+  {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    PodTemplateWithName that = (PodTemplateWithName) o;
+    return Objects.equals(name, that.name) &&
+        Objects.equals(podTemplate, that.podTemplate);
+  }
+
+  @Override
+  public int hashCode()
+  {
+    return Objects.hash(name, podTemplate);
+  }
+
+  @Override
+  public String toString()
+  {
+    return "PodTemplateWithName{" +
+        "name='" + name + '\'' +
+        ", podTemplate=" + podTemplate +
+        '}';
+  }
+}

extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/common/PodTemplateWithNameTest.java

@@ -0,0 +1,60 @@
+/*
+ * 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.common;
+
+import io.fabric8.kubernetes.api.model.PodTemplate;
+import io.fabric8.kubernetes.api.model.PodTemplateBuilder;
+import org.apache.druid.k8s.overlord.taskadapter.PodTemplateWithName;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+
+public class PodTemplateWithNameTest
+{
+  @Test
+  void testEqualityToMakeCoverageHappy()
+  {
+    PodTemplateWithName podTemplateWithName = new PodTemplateWithName(
+        "name",
+        new PodTemplateBuilder().build()
+    );
+    PodTemplateWithName podTemplateWithName2 = podTemplateWithName;
+
+    assertEquals(podTemplateWithName, podTemplateWithName2);
+    assertNotEquals(podTemplateWithName, null);
+    assertNotEquals(podTemplateWithName, "string");
+    assertEquals(podTemplateWithName.hashCode(), podTemplateWithName2.hashCode());
+  }
+
+  @Test
+  void testGettersToMakeCoverageHappy()
+  {
+    String name = "name";
+    PodTemplate podTemplate = new PodTemplateBuilder().build();
+    PodTemplateWithName podTemplateWithName = new PodTemplateWithName(
+        name,
+        podTemplate
+    );
+
+    assertEquals(name, podTemplateWithName.getName());
+    assertEquals(podTemplate, podTemplateWithName.getPodTemplate());
+  }
+}

extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/execution/SelectorBasedPodTemplateSelectStrategyTest.java

@@ -27,6 +27,7 @@
 import io.fabric8.kubernetes.api.model.PodTemplate;
 import org.apache.druid.indexing.common.task.NoopTask;
 import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.k8s.overlord.taskadapter.PodTemplateWithName;
 import org.apache.druid.segment.TestHelper;
 import org.junit.Assert;
 import org.junit.Before;
@@ -97,7 +98,10 @@ public void testGetPodTemplate_ForTask_emptySelectorsFallbackToBaseTemplate()
     List<Selector> emptySelectors = Collections.emptyList();
     SelectorBasedPodTemplateSelectStrategy strategy = new SelectorBasedPodTemplateSelectStrategy(emptySelectors);
     Task task = NoopTask.create();
-    Assert.assertEquals("base", strategy.getPodTemplateForTask(task, templates).getMetadata().getName());
+    PodTemplateWithName podTemplateWithName = strategy.getPodTemplateForTask(task, templates);
+    Assert.assertEquals("base", podTemplateWithName.getName());
+    Assert.assertEquals("base", podTemplateWithName.getPodTemplate().getMetadata().getName());
+
   }
 
   @Test
@@ -107,7 +111,9 @@ public void testGetPodTemplate_ForTask_noMatchSelectorsFallbackToBaseTemplateIfN
     List<Selector> selectors = Collections.singletonList(noMatchSelector);
     SelectorBasedPodTemplateSelectStrategy strategy = new SelectorBasedPodTemplateSelectStrategy(selectors);
     Task task = NoopTask.create();
-    Assert.assertEquals("base", strategy.getPodTemplateForTask(task, templates).getMetadata().getName());
+    PodTemplateWithName podTemplateWithName = strategy.getPodTemplateForTask(task, templates);
+    Assert.assertEquals("base", podTemplateWithName.getName());
+    Assert.assertEquals("base", podTemplateWithName.getPodTemplate().getMetadata().getName());
   }
 
   @Test
@@ -124,7 +130,9 @@ public void testGetPodTemplate_ForTask_withMatchSelectors()
     );
     SelectorBasedPodTemplateSelectStrategy strategy = new SelectorBasedPodTemplateSelectStrategy(selectors);
     Task task = NoopTask.create();
-    Assert.assertEquals("match", strategy.getPodTemplateForTask(task, templates).getMetadata().getName());
+    PodTemplateWithName podTemplateWithName = strategy.getPodTemplateForTask(task, templates);
+    Assert.assertEquals("match", podTemplateWithName.getName());
+    Assert.assertEquals("match", podTemplateWithName.getPodTemplate().getMetadata().getName());
   }
 
   @Test

extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/taskadapter/PodTemplateTaskAdapterTest.java

@@ -172,7 +172,7 @@ public void test_fromTask_withBasePodTemplateInRuntimeProperites() throws IOExce
         null
     );
     Job actual = adapter.fromTask(task);
-    Job expected = K8sTestUtils.fileToResource("expectedNoopJob.yaml", Job.class);
+    Job expected = K8sTestUtils.fileToResource("expectedNoopJobBase.yaml", Job.class);
 
     assertJobSpecsEqual(actual, expected);
   }
@@ -216,7 +216,7 @@ public void test_fromTask_withBasePodTemplateInRuntimeProperites_andTlsEnabled()
     );
 
     Job actual = adapter.fromTask(task);
-    Job expected = K8sTestUtils.fileToResource("expectedNoopJobTlsEnabled.yaml", Job.class);
+    Job expected = K8sTestUtils.fileToResource("expectedNoopJobTlsEnabledBase.yaml", Job.class);
 
     assertJobSpecsEqual(actual, expected);
   }

extensions-contrib/kubernetes-overlord-extensions/src/test/resources/expectedNoopJob.yaml

@@ -8,11 +8,13 @@ metadata:
     druid.task.type: "noop"
     druid.task.group.id: "id"
     druid.task.datasource: "datasource"
+
   annotations:
     task.id: "id"
     task.type: "noop"
     task.group.id: "id"
     task.datasource: "datasource"
+    task.jobTemplate: noop
 spec:
   activeDeadlineSeconds: 14400
   backoffLimit: 0
@@ -32,6 +34,7 @@ spec:
         task.type: "noop"
         task.group.id: "id"
         task.datasource: "datasource"
+        task.jobTemplate: noop
     spec:
       containers:
         - command:

extensions-contrib/kubernetes-overlord-extensions/src/test/resources/expectedNoopJobBase.yaml

@@ -0,0 +1,55 @@
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: "id-3e70afe5cd823dfc7dd308eea616426b"
+  labels:
+    druid.k8s.peons: "true"
+    druid.task.id: "id"
+    druid.task.type: "noop"
+    druid.task.group.id: "id"
+    druid.task.datasource: "datasource"
+
+  annotations:
+    task.id: "id"
+    task.type: "noop"
+    task.group.id: "id"
+    task.datasource: "datasource"
+    task.jobTemplate: base
+spec:
+  activeDeadlineSeconds: 14400
+  backoffLimit: 0
+  ttlSecondsAfterFinished: 172800
+  template:
+    metadata:
+      labels:
+        druid.k8s.peons: "true"
+        druid.task.id: "id"
+        druid.task.type: "noop"
+        druid.task.group.id: "id"
+        druid.task.datasource: "datasource"
+      annotations:
+        task: "H4sIAAAAAAAAAD2MvQ4CIRCE32VqijsTG1qLi7W+wArEbHICrmC8EN7dJf40k/lmJtNQthxgEVPKMGCvXsXgKqnm4x89FTqlKm6MBzw+YCA1nvmm8W4/TQYuxRJeBbZ17cJ3ZhvoSbzShVcu2zLOf9cS7pUl+ANlclrCzr2/AQUK0FqZAAAA"
+        tls.enabled: "false"
+        task.id: "id"
+        task.type: "noop"
+        task.group.id: "id"
+        task.datasource: "datasource"
+        task.jobTemplate: base
+    spec:
+      containers:
+        - command:
+            - sleep
+            - "3600"
+          env:
+            - name: "TASK_DIR"
+              value: "/tmp"
+            - name: "TASK_ID"
+              value: "id"
+            - name: "LOAD_BROADCAST_SEGMENTS"
+              value: "false"
+            - name: "TASK_JSON"
+              valueFrom:
+                fieldRef:
+                  fieldPath: "metadata.annotations['task']"
+          image: one
+          name: primary