Add support for dynamic kubernetes labels on pod and service (#3701)

apache · Jul 9, 2021 · 6aa6ae6 · 6aa6ae6
1 parent bded8f5
commit 6aa6ae6
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 22 deletions.
diff --git a/heron/schedulers/src/java/org/apache/heron/scheduler/kubernetes/KubernetesContext.java b/heron/schedulers/src/java/org/apache/heron/scheduler/kubernetes/KubernetesContext.java
@@ -92,6 +92,11 @@ public enum KubernetesResourceRequestMode {
       "heron.kubernetes.pod.annotation.";
   public static final String HERON_KUBERNETES_SERVICE_ANNOTATION =
       "heron.kubernetes.service.annotation.";
+  public static final String HERON_KUBERNETES_POD_LABEL =
+          "heron.kubernetes.pod.label.";
+  public static final String HERON_KUBERNETES_SERVICE_LABEL =
+          "heron.kubernetes.service.label.";
+
 
   private KubernetesContext() {
   }
@@ -162,6 +167,22 @@ static String getContainerVolumeMountPath(Config config) {
     return config.getStringValue(HERON_KUBERNETES_CONTAINER_VOLUME_MOUNT_PATH);
   }
 
+  public static Map<String, String> getPodLabels(Config config) {
+    return getConfigItemsByPrefix(config, HERON_KUBERNETES_POD_LABEL);
+  }
+
+  public static Map<String, String> getServiceLabels(Config config) {
+    return getConfigItemsByPrefix(config, HERON_KUBERNETES_SERVICE_LABEL);
+  }
+
+  public static Map<String, String> getPodAnnotations(Config config) {
+    return getConfigItemsByPrefix(config, HERON_KUBERNETES_POD_ANNOTATION);
+  }
+
+  public static Map<String, String> getServiceAnnotations(Config config) {
+    return getConfigItemsByPrefix(config, HERON_KUBERNETES_SERVICE_ANNOTATION);
+  }
+
   static Set<String> getConfigKeys(Config config, String keyPrefix) {
     Set<String> annotations = new HashSet<>();
     for (String s : config.getKeySet()) {
@@ -172,26 +193,14 @@ static Set<String> getConfigKeys(Config config, String keyPrefix) {
     return annotations;
   }
 
-  public static Map<String, String> getPodAnnotations(Config config) {
-    final Map<String, String> annotations = new HashMap<>();
-    final Set<String> keys = getConfigKeys(config, HERON_KUBERNETES_POD_ANNOTATION);
-    for (String s : keys) {
-      String value = config.getStringValue(s);
-      annotations.put(s.replaceFirst(KubernetesContext.HERON_KUBERNETES_POD_ANNOTATION,
-              ""), value);
-    }
-    return annotations;
-  }
-
-  public static Map<String, String> getServiceAnnotations(Config config) {
-    final Map<String, String> annotations = new HashMap<>();
-    final Set<String> keys = getConfigKeys(config, HERON_KUBERNETES_SERVICE_ANNOTATION);
+  private static Map<String, String> getConfigItemsByPrefix(Config config, String keyPrefix) {
+    final Map<String, String> results = new HashMap<>();
+    final Set<String> keys = getConfigKeys(config, keyPrefix);
     for (String s : keys) {
       String value = config.getStringValue(s);
-      annotations.put(s.replaceFirst(KubernetesContext.HERON_KUBERNETES_SERVICE_ANNOTATION,
-              ""), value);
+      results.put(s.replaceFirst(keyPrefix, ""), value);
     }
-    return annotations;
+    return results;
   }
 
   public static boolean hasContainerVolume(Config config) {

diff --git a/heron/schedulers/src/java/org/apache/heron/scheduler/kubernetes/V1Controller.java b/heron/schedulers/src/java/org/apache/heron/scheduler/kubernetes/V1Controller.java
@@ -327,12 +327,13 @@ private V1Service createTopologyService() {
     final V1ObjectMeta objectMeta = new V1ObjectMeta();
     objectMeta.name(topologyName);
     objectMeta.annotations(getServiceAnnotations());
+    objectMeta.setLabels(getServiceLabels());
     service.setMetadata(objectMeta);
 
     // create the headless service
     final V1ServiceSpec serviceSpec = new V1ServiceSpec();
     serviceSpec.clusterIP("None");
-    serviceSpec.setSelector(getMatchLabels(topologyName));
+    serviceSpec.setSelector(getPodMatchLabels(topologyName));
 
     service.setSpec(serviceSpec);
 
@@ -363,14 +364,14 @@ private V1StatefulSet createStatefulSet(Resource containerResource, int numberOf
     // add selector match labels "app=heron" and "topology=topology-name"
     // so the we know which pods to manage
     final V1LabelSelector selector = new V1LabelSelector();
-    selector.matchLabels(getMatchLabels(topologyName));
+    selector.matchLabels(getPodMatchLabels(topologyName));
     statefulSetSpec.selector(selector);
 
     // create a pod template
     final V1PodTemplateSpec podTemplateSpec = new V1PodTemplateSpec();
 
     // set up pod meta
-    final V1ObjectMeta templateMetaData = new V1ObjectMeta().labels(getLabels(topologyName));
+    final V1ObjectMeta templateMetaData = new V1ObjectMeta().labels(getPodLabels(topologyName));
     Map<String, String> annotations = new HashMap<>();
     annotations.putAll(getPodAnnotations());
     annotations.putAll(getPrometheusAnnotations());
@@ -408,20 +409,25 @@ private Map<String, String> getPrometheusAnnotations() {
     return annotations;
   }
 
-  private Map<String, String> getMatchLabels(String topologyName) {
+  private Map<String, String> getPodMatchLabels(String topologyName) {
     final Map<String, String> labels = new HashMap<>();
     labels.put(KubernetesConstants.LABEL_APP, KubernetesConstants.LABEL_APP_VALUE);
     labels.put(KubernetesConstants.LABEL_TOPOLOGY, topologyName);
     return labels;
   }
 
-  private Map<String, String> getLabels(String topologyName) {
+  private Map<String, String> getPodLabels(String topologyName) {
     final Map<String, String> labels = new HashMap<>();
     labels.put(KubernetesConstants.LABEL_APP, KubernetesConstants.LABEL_APP_VALUE);
     labels.put(KubernetesConstants.LABEL_TOPOLOGY, topologyName);
+    labels.putAll(KubernetesContext.getPodLabels(getConfiguration()));
     return labels;
   }
 
+  private Map<String, String> getServiceLabels() {
+    return KubernetesContext.getServiceLabels(getConfiguration());
+  }
+
   private V1PodSpec getPodSpec(List<String> executorCommand, Resource resource,
       int numberOfInstances) {
     final V1PodSpec podSpec = new V1PodSpec();