Repository: mesos
Updated Branches:
  refs/heads/master 11105f7c0 -> 01ee1e504


Enabled label decorator to override.

Review: https://reviews.apache.org/r/30961


Project: http://git-wip-us.apache.org/repos/asf/mesos/repo
Commit: http://git-wip-us.apache.org/repos/asf/mesos/commit/a090e9dc
Tree: http://git-wip-us.apache.org/repos/asf/mesos/tree/a090e9dc
Diff: http://git-wip-us.apache.org/repos/asf/mesos/diff/a090e9dc

Branch: refs/heads/master
Commit: a090e9dc7698ecfe5da51826c1aa2e5163966272
Parents: 11105f7
Author: Niklas Nielsen <[email protected]>
Authored: Mon Apr 20 14:36:14 2015 -0700
Committer: Adam B <[email protected]>
Committed: Mon Apr 20 14:36:14 2015 -0700

----------------------------------------------------------------------
 include/mesos/hook.hpp            |  6 +++---
 src/examples/test_hook_module.cpp | 16 +++++++++++++---
 src/hook/manager.cpp              | 15 +++++++++++----
 src/master/master.cpp             |  4 ++--
 src/tests/hook_tests.cpp          | 25 +++++++++++++++++--------
 5 files changed, 46 insertions(+), 20 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/mesos/blob/a090e9dc/include/mesos/hook.hpp
----------------------------------------------------------------------
diff --git a/include/mesos/hook.hpp b/include/mesos/hook.hpp
index 9ae8b94..f2b8259 100644
--- a/include/mesos/hook.hpp
+++ b/include/mesos/hook.hpp
@@ -34,9 +34,9 @@ public:
   virtual ~Hook() {};
 
   // This label decorator hook is called from within master during
-  // the launchTask routine.  A module implementing the hook creates
-  // and returns a set of labels.  These labels are then merged with
-  // the task labels and passed on to the slave/executor.
+  // the launchTask routine. A module implementing the hook creates
+  // and returns a set of labels. These labels overwrite the existing
+  // labels on the task info.
   virtual Result<Labels> masterLaunchTaskLabelDecorator(
       const TaskInfo& taskInfo,
       const FrameworkInfo& frameworkInfo,

http://git-wip-us.apache.org/repos/asf/mesos/blob/a090e9dc/src/examples/test_hook_module.cpp
----------------------------------------------------------------------
diff --git a/src/examples/test_hook_module.cpp 
b/src/examples/test_hook_module.cpp
index 2f2da1c..4fc7013 100644
--- a/src/examples/test_hook_module.cpp
+++ b/src/examples/test_hook_module.cpp
@@ -32,6 +32,7 @@ using namespace mesos;
 // tests/hook_tests.cpp.
 const char* testLabelKey = "MESOS_Test_Label";
 const char* testLabelValue = "ApacheMesos";
+const char* testRemoveLabelKey = "MESOS_Test_Remove_Label";
 
 
 class TestHook : public Hook
@@ -45,9 +46,18 @@ public:
     LOG(INFO) << "Executing 'masterLaunchTaskLabelDecorator' hook";
 
     Labels labels;
-    Label *label = labels.add_labels();
-    label->set_key(testLabelKey);
-    label->set_value(testLabelValue);
+
+    // Set one known label.
+    Label* newLabel = labels.add_labels();
+    newLabel->set_key(testLabelKey);
+    newLabel->set_value(testLabelValue);
+
+    // Remove the 'testRemoveLabelKey' label which was set by the test.
+    foreach (const Label& oldLabel, taskInfo.labels().labels()) {
+      if (oldLabel.key() != testRemoveLabelKey) {
+        labels.add_labels()->CopyFrom(oldLabel);
+      }
+    }
 
     return labels;
   }

http://git-wip-us.apache.org/repos/asf/mesos/blob/a090e9dc/src/hook/manager.cpp
----------------------------------------------------------------------
diff --git a/src/hook/manager.cpp b/src/hook/manager.cpp
index 7a4cb09..0d79ba2 100644
--- a/src/hook/manager.cpp
+++ b/src/hook/manager.cpp
@@ -96,20 +96,27 @@ Labels HookManager::masterLaunchTaskLabelDecorator(
     const SlaveInfo& slaveInfo)
 {
   Lock lock(&mutex);
-  Labels labels;
+
+  // We need a mutable copy of the task info and set the new
+  // labels after each hook invocation. Otherwise, the last hook
+  // will be the only effective hook setting the labels.
+  TaskInfo taskInfo_ = taskInfo;
 
   foreachpair (const string& name, Hook* hook, availableHooks) {
     const Result<Labels>& result =
-      hook->masterLaunchTaskLabelDecorator(taskInfo, frameworkInfo, slaveInfo);
+      hook->masterLaunchTaskLabelDecorator(taskInfo_, frameworkInfo, 
slaveInfo);
+
+    // NOTE: If the hook returns None(), the task labels won't be
+    // changed.
     if (result.isSome()) {
-      labels.MergeFrom(result.get());
+      taskInfo_.mutable_labels()->CopyFrom(result.get());
     } else if (result.isError()) {
       LOG(WARNING) << "Master label decorator hook failed for module '"
                    << name << "': " << result.error();
     }
   }
 
-  return labels;
+  return taskInfo_.labels();
 }
 
 

http://git-wip-us.apache.org/repos/asf/mesos/blob/a090e9dc/src/master/master.cpp
----------------------------------------------------------------------
diff --git a/src/master/master.cpp b/src/master/master.cpp
index e30b951..b1093bb 100644
--- a/src/master/master.cpp
+++ b/src/master/master.cpp
@@ -2619,8 +2619,8 @@ void Master::_accept(
             message.set_pid(framework->pid);
             message.mutable_task()->MergeFrom(task);
 
-            // Merge labels retrieved from label-decorator hooks.
-            message.mutable_task()->mutable_labels()->MergeFrom(
+            // Set labels retrieved from label-decorator hooks.
+            message.mutable_task()->mutable_labels()->CopyFrom(
                 HookManager::masterLaunchTaskLabelDecorator(
                     task,
                     framework->info,

http://git-wip-us.apache.org/repos/asf/mesos/blob/a090e9dc/src/tests/hook_tests.cpp
----------------------------------------------------------------------
diff --git a/src/tests/hook_tests.cpp b/src/tests/hook_tests.cpp
index bb9de25..fdfcd38 100644
--- a/src/tests/hook_tests.cpp
+++ b/src/tests/hook_tests.cpp
@@ -74,6 +74,8 @@ const char* HOOK_MODULE_NAME = "org_apache_mesos_TestHook";
 // examples/test_hook_module.cpp.
 const char* testLabelKey = "MESOS_Test_Label";
 const char* testLabelValue = "ApacheMesos";
+const char* testRemoveLabelKey = "MESOS_Test_Remove_Label";
+const char* testRemoveLabelValue = "FooBar";
 const char* testEnvironmentVariableName = "MESOS_TEST_ENVIRONMENT_VARIABLE";
 
 class HookTest : public MesosTest
@@ -162,6 +164,12 @@ TEST_F(HookTest, VerifyMasterLaunchTaskHook)
   task.mutable_resources()->CopyFrom(offers.get()[0].resources());
   task.mutable_command()->CopyFrom(command);
 
+  // Add label which will be removed by the hook.
+  Labels* labels = task.mutable_labels();
+  Label* label = labels->add_labels();
+  label->set_key(testRemoveLabelKey);
+  label->set_value(testRemoveLabelValue);
+
   vector<TaskInfo> tasks;
   tasks.push_back(task);
 
@@ -175,14 +183,15 @@ TEST_F(HookTest, VerifyMasterLaunchTaskHook)
   AWAIT_READY(taskInfo);
 
   // At launchTasks, the label decorator hook inside should have been
-  // executed and we should see the labels now.
-  Option<string> labelValue;
-  foreach (const Label& label, taskInfo.get().labels().labels()) {
-    if (label.key() == testLabelKey) {
-      labelValue = label.value();
-    }
-  }
-  EXPECT_SOME_EQ(testLabelValue, labelValue);
+  // executed and we should see the labels now. Also, verify that the
+  // hook module has stripped the first 'testRemoveLabelKey' label.
+  // We do this by ensuring that only one label is present and that it
+  // is the new 'testLabelKey' label.
+  const Labels &labels_ = taskInfo.get().labels();
+  ASSERT_EQ(1, labels_.labels_size());
+
+  EXPECT_EQ(labels_.labels().Get(0).key(), testLabelKey);
+  EXPECT_EQ(labels_.labels().Get(0).value(), testLabelValue);
 
   driver.stop();
   driver.join();

Reply via email to