From 91cdc7e91afa774ca7eb5a04232cf8591b02f570 Mon Sep 17 00:00:00 2001
From: Kan Wu <optimuswu8685@gmail.com>
Date: Wed, 14 Dec 2016 04:32:19 -0800
Subject: [PATCH] add a docker container label with mesos task id (#4492)

---
 .../scala/mesosphere/mesos/TaskBuilder.scala  | 15 ++++++---
 .../mesosphere/mesos/TaskBuilderTest.scala    | 32 +++++++++++++++++--
 2 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/src/main/scala/mesosphere/mesos/TaskBuilder.scala b/src/main/scala/mesosphere/mesos/TaskBuilder.scala
index d4a76a15efa..355539c50ad 100644
--- a/src/main/scala/mesosphere/mesos/TaskBuilder.scala
+++ b/src/main/scala/mesosphere/mesos/TaskBuilder.scala
@@ -56,7 +56,7 @@ class TaskBuilder(
 
     volumeMatchOpt.foreach(_.persistentVolumeResources.foreach(builder.addResources))
 
-    val containerProto = computeContainerInfo(resourceMatch.hostPorts)
+    val containerProto = computeContainerInfo(resourceMatch.hostPorts, taskId)
     val envPrefix: Option[String] = config.envVarsPrefix.get
 
     executor match {
@@ -155,7 +155,7 @@ class TaskBuilder(
     discoveryInfoBuilder.build
   }
 
-  protected def computeContainerInfo(hostPorts: Seq[Option[Int]]): Option[ContainerInfo] = {
+  protected def computeContainerInfo(hostPorts: Seq[Option[Int]], taskId: Task.Id): Option[ContainerInfo] = {
     if (runSpec.container.isEmpty && runSpec.ipAddress.isEmpty) {
       None
     } else {
@@ -185,12 +185,17 @@ class TaskBuilder(
         // TODO(portMappings)
         // TODO(nfnt): Other containers might also support port mappings in the future.
         // If that is the case, a more general way than the one below needs to be implemented.
-        val containerWithPortMappings = c match {
-          case docker: Container.Docker => docker.copy(portMappings = boundPortMappings)
+        val updatedContainer = c match {
+          case docker: Container.Docker =>
+            docker.copy(
+              portMappings = boundPortMappings,
+              parameters = docker.parameters :+
+                new mesosphere.marathon.state.Parameter("label", s"MESOS_TASK_ID=${taskId.mesosTaskId.getValue}")
+            )
           case _ => c
         }
 
-        builder.mergeFrom(ContainerSerializer.toMesos(containerWithPortMappings))
+        builder.mergeFrom(ContainerSerializer.toMesos(updatedContainer))
       }
 
       // Set NetworkInfo if necessary
diff --git a/src/test/scala/mesosphere/mesos/TaskBuilderTest.scala b/src/test/scala/mesosphere/mesos/TaskBuilderTest.scala
index 4daed2b0979..645716fc116 100644
--- a/src/test/scala/mesosphere/mesos/TaskBuilderTest.scala
+++ b/src/test/scala/mesosphere/mesos/TaskBuilderTest.scala
@@ -1174,6 +1174,33 @@ class TaskBuilderTest extends MarathonSpec with Matchers {
     assert(containerPort == hostPort)
   }
 
+  test("DockerContainerWithMesosTaskIdLabel") {
+    val offer = MarathonTestHelper.makeBasicOfferWithRole(
+      cpus = 1.0, mem = 128.0, disk = 1000.0, beginPort = 31000, endPort = 31010, role = ResourceRole.Unreserved
+    )
+      .addResources(RangesResource(Resource.PORTS, Seq(protos.Range(33000, 34000)), "marathon"))
+      .build
+
+    val task: Option[(MesosProtos.TaskInfo, _)] = buildIfMatches(
+      offer, AppDefinition(
+        id = "testApp".toPath,
+        resources = Resources(cpus = 1.0, mem = 64.0, disk = 1.0),
+        executor = "//cmd",
+        container = Some(Docker(
+          image = "busybox"
+        ))
+      ), None, None, None,
+      _ => Task.Id("mesos_task_id")
+    )
+    assert(task.isDefined, "expected task to match offer")
+    val (taskInfo, _) = task.get
+
+    assert(taskInfo.getContainer.getDocker.getParametersList.size == 1, s"expected 1 parameter, but ${taskInfo.getContainer.getDocker.getParametersList.size}")
+    val param = taskInfo.getContainer.getDocker.getParametersList.get(0)
+    assert(param.getKey == "label", "expected docker having a parameter key: label")
+    assert(param.getValue == "MESOS_TASK_ID=mesos_task_id", s"expected docker having a parameter value for key 'label': MESOS_TASK_ID=mesos_task_id but ${param.getValue }")
+  }
+
   test("BuildIfMatchesWithRackIdConstraint") {
     val offer = MarathonTestHelper.makeBasicOffer(1.0, 128.0, 31000, 32000)
       .addAttributes(TextAttribute("rackid", "1"))
@@ -1781,10 +1808,11 @@ class TaskBuilderTest extends MarathonSpec with Matchers {
     app: AppDefinition,
     mesosRole: Option[String] = None,
     acceptedResourceRoles: Option[Set[String]] = None,
-    envVarsPrefix: Option[String] = None): Option[(MesosProtos.TaskInfo, NetworkInfo)] = {
+    envVarsPrefix: Option[String] = None,
+    newTaskId: PathId => Task.Id = s => Task.Id.forRunSpec(s)): Option[(MesosProtos.TaskInfo, NetworkInfo)] = {
     val builder = new TaskBuilder(
       app,
-      s => Task.Id.forRunSpec(s),
+      newTaskId,
       MarathonTestHelper.defaultConfig(
         mesosRole = mesosRole,
         acceptedResourceRoles = acceptedResourceRoles,