Merge pull request #545 from dylanmei/mesos_task_labels

Mesos labels can be added to executor tasks

Author: frankscholten

docs/index.md

@@ -158,6 +158,9 @@ Usage: (Options preceded by an asterisk are required) [options]
     --executorName
        The name given to the executor task.
        Default: elasticsearch-executor
+    --executorLabels
+       One or more labels given to the executor task. Example: 'environment=prod bananas=apples'
+       Default: <empty string>
     --externalVolumeDriver
        Use external volume storage driver. By default, nodes will use volumes on
        host.

scheduler/src/main/java/org/apache/mesos/elasticsearch/scheduler/Configuration.java

@@ -19,6 +19,8 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
+import java.util.HashMap;
 import java.util.stream.Collectors;
 
 import static java.util.Arrays.asList;
@@ -36,6 +38,7 @@ public class Configuration {
     public static final String WEB_UI_PORT = "--webUiPort";
     public static final String FRAMEWORK_NAME = "--frameworkName";
     public static final String EXECUTOR_NAME = "--executorName";
+    public static final String EXECUTOR_LABELS = "--executorLabels";
     public static final String DATA_DIR = "--dataDir";
     public static final String DEFAULT_HOST_DATA_DIR = "/var/lib/mesos/slave/elasticsearch";
     // DCOS Certification requirement 01
@@ -85,6 +88,9 @@ public class Configuration {
     private String frameworkName = "elasticsearch";
     @Parameter(names = {EXECUTOR_NAME}, description = "The name given to the executor task.", validateWith = CLIValidators.NotEmptyString.class)
     private String executorName = "elasticsearch-executor";
+    @Parameter(names = {EXECUTOR_LABELS}, description = "One or more labels given to the executor task." +
+        "E.g. 'environment=prod bananas=apples'", variableArity = true)
+    private List<String> executorLabels = new ArrayList<>();
     @Parameter(names = {DATA_DIR}, description = "The host data directory used by Docker volumes in the executors. [DOCKER MODE ONLY]")
     private String dataDir = DEFAULT_HOST_DATA_DIR;
     @Parameter(names = {FRAMEWORK_FAILOVER_TIMEOUT}, description = "The time before Mesos kills a scheduler and tasks if it has not recovered (ms).", validateValueWith = CLIValidators.PositiveDouble.class)
@@ -176,6 +182,18 @@ public String getTaskName() {
         return executorName;
     }
 
+    public Map<String, String> getTaskLabels() {
+        HashMap<String, String> map = new HashMap<>();
+        for (String keyValue : executorLabels) {
+            String[] kvp = keyValue.split("=", 2);
+            if (kvp.length == 2) {
+              map.put(kvp[0], kvp[1]);
+            }
+        }
+
+        return map;
+    }
+
     public String getDataDir() {
         return dataDir;
     }

scheduler/src/main/java/org/apache/mesos/elasticsearch/scheduler/TaskInfoFactory.java

@@ -19,6 +19,7 @@
 import java.time.ZoneOffset;
 import java.time.ZonedDateTime;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.stream.Collectors;
@@ -69,6 +70,7 @@ private Protos.TaskInfo buildNativeTask(Protos.Offer offer, Configuration config
         final List<Integer> ports = getPorts(offer, configuration);
         final List<Protos.Resource> resources = getResources(configuration, ports);
         final Protos.DiscoveryInfo discovery = getDiscovery(ports, configuration);
+        final Protos.Labels labels = getLabels(configuration);
 
         final String hostAddress = resolveHostAddress(offer, ports);
 
@@ -83,6 +85,7 @@ private Protos.TaskInfo buildNativeTask(Protos.Offer offer, Configuration config
                 .setSlaveId(offer.getSlaveId())
                 .addAllResources(resources)
                 .setDiscovery(discovery)
+                .setLabels(labels)
                 .setCommand(nativeCommand(configuration, args, elasticSearchNodeId))
                 .build();
     }
@@ -91,6 +94,7 @@ private Protos.TaskInfo buildDockerTask(Protos.Offer offer, Configuration config
         final List<Integer> ports = getPorts(offer, configuration);
         final List<Protos.Resource> resources = getResources(configuration, ports);
         final Protos.DiscoveryInfo discovery = getDiscovery(ports, configuration);
+        final Protos.Labels labels = getLabels(configuration);
 
         final String hostAddress = resolveHostAddress(offer, ports);
 
@@ -107,6 +111,7 @@ private Protos.TaskInfo buildDockerTask(Protos.Offer offer, Configuration config
                 .setSlaveId(offer.getSlaveId())
                 .addAllResources(resources)
                 .setDiscovery(discovery)
+                .setLabels(labels)
                 .setCommand(dockerCommand(configuration, args, elasticSearchNodeId))
                 .setContainer(containerInfo)
                 .build();
@@ -151,6 +156,21 @@ private Protos.DiscoveryInfo getDiscovery(List<Integer> ports, Configuration con
         return discovery.build();
     }
 
+    private Protos.Labels getLabels(Configuration configuration) {
+      Protos.Labels.Builder builder = Protos.Labels.newBuilder();
+      Map<String, String> labels = configuration.getTaskLabels();
+      for (Map.Entry<String, String> kvp : labels.entrySet()) {
+        Protos.Label label = Protos.Label.newBuilder()
+          .setKey(kvp.getKey())
+          .setValue(kvp.getValue())
+          .build();
+
+        builder.addLabels(label);
+      }
+
+      return builder.build();
+    }
+
     private Protos.ContainerInfo getContainer(Configuration configuration, Protos.TaskID taskID, Long elasticSearchNodeId, Protos.SlaveID slaveID) {
         final Protos.Environment environment = Protos.Environment.newBuilder().addAllVariables(new ExecutorEnvironmentalVariables(configuration, elasticSearchNodeId).getList()).build();
         final Protos.ContainerInfo.DockerInfo.Builder dockerInfo = Protos.ContainerInfo.DockerInfo.newBuilder()

scheduler/src/test/java/org/apache/mesos/elasticsearch/scheduler/ConfigurationTest.java

@@ -11,6 +11,7 @@
 import java.net.UnknownHostException;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 
 import static org.junit.Assert.*;
 
@@ -78,4 +79,20 @@ public void shouldCreateVolumeName() {
         Configuration configuration = new Configuration(ZookeeperCLIParameter.ZOOKEEPER_MESOS_URL, "aa", Configuration.FRAMEWORK_NAME, "test");
         assertEquals("test0data", configuration.dataVolumeName(0L));
     }
-}
+
+    @Test
+    public void shouldCreateTaskLabels() {
+        Configuration configuration = new Configuration(
+            ZookeeperCLIParameter.ZOOKEEPER_MESOS_URL, "aa", Configuration.EXECUTOR_LABELS,
+              "foo=bar",
+              "incomplete",
+              "empty=",
+              "separator=values=are=joined");
+        Map<String, String> labels = configuration.getTaskLabels();
+
+        assertEquals("bar", labels.get("foo"));
+        assertFalse(labels.containsKey("incomplete"));
+        assertEquals("", labels.get("empty"));
+        assertEquals("values=are=joined", labels.get("separator"));
+    }
+}