[mesos] 03/04: Added ability to suppress a subset of roles in the scheduler driver.

[bm...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

bmahler pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/mesos.git

commit 57c2b7098ee523f3f0f647c6b334cce37266fe2a
Author: Andrei Sekretenko <as...@mesosphere.io>
AuthorDate: Wed Jul 3 18:47:53 2019 -0400

    Added ability to suppress a subset of roles in the scheduler driver.
    
    This patch adds to the scheduler driver a 'suppressOffers(roles)'
    method, which sends the SUPPRESS call for these roles and adds them to
    the suppressed roles set.
    
    Review: https://reviews.apache.org/r/70983/
---
 include/mesos/scheduler.hpp |  9 +++++++++
 src/sched/sched.cpp         | 44 ++++++++++++++++++++++++++++++++++++++------
 2 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/include/mesos/scheduler.hpp b/include/mesos/scheduler.hpp
index cc3a278..61cc846 100644
--- a/include/mesos/scheduler.hpp
+++ b/include/mesos/scheduler.hpp
@@ -312,6 +312,13 @@ public:
   // re-registration.
   virtual Status suppressOffers() = 0;
 
+  // Adds the roles to the suppressed set. If the framework is not connected
+  // to the master, an up-to-date set of suppressed roles will be sent to
+  // the master during re-registration.
+  //
+  // NOTE: If 'roles' is empty, this method does nothing.
+  virtual Status suppressOffers(const std::vector<std::string>& roles) = 0;
+
   // Acknowledges the status update. This should only be called
   // once the status update is processed durably by the scheduler.
   // Not that explicit acknowledgements must be requested via the
@@ -502,6 +509,8 @@ public:
 
   Status suppressOffers() override;
 
+  Status suppressOffers(const std::vector<std::string>& roles) override;
+
   Status acknowledgeStatusUpdate(
       const TaskStatus& status) override;
 
diff --git a/src/sched/sched.cpp b/src/sched/sched.cpp
index 47b87ad..768ce7d 100644
--- a/src/sched/sched.cpp
+++ b/src/sched/sched.cpp
@@ -1475,10 +1475,16 @@ protected:
     send(master->pid(), call);
   }
 
-  void suppressOffers()
+  void suppressOffers(const vector<string>& roles)
   {
-    suppressedRoles =
-      std::set<string>(framework.roles().begin(), framework.roles().end());
+    if (roles.empty()) {
+      suppressedRoles =
+        std::set<string>(framework.roles().begin(), framework.roles().end());
+    } else {
+      for (const string& role : roles) {
+        suppressedRoles.emplace(role);
+      }
+    }
 
     if (!connected) {
       VLOG(1) << "Ignoring SUPPRESS as master is disconnected;"
@@ -1489,14 +1495,20 @@ protected:
       return;
     }
 
-    VLOG(2) << "Sending SUPPRESS for all roles";
-
     Call call;
 
     CHECK(framework.has_id());
     call.mutable_framework_id()->CopyFrom(framework.id());
     call.set_type(Call::SUPPRESS);
 
+    if (roles.empty()) {
+      VLOG(2) << "Sending SUPPRESS for all roles";
+    } else {
+      VLOG(2) << "Sending SUPPRESS for roles: " << stringify(roles);
+      *call.mutable_suppress()->mutable_roles() =
+        RepeatedPtrField<string>(roles.begin(), roles.end());
+    }
+
     CHECK_SOME(master);
     send(master->pid(), call);
   }
@@ -2375,7 +2387,27 @@ Status MesosSchedulerDriver::suppressOffers()
 
     CHECK(process != nullptr);
 
-    dispatch(process, &SchedulerProcess::suppressOffers);
+    dispatch(process, &SchedulerProcess::suppressOffers, vector<string>());
+
+    return status;
+  }
+}
+
+
+Status MesosSchedulerDriver::suppressOffers(const vector<string>& roles)
+{
+  if (roles.empty()) {
+    return status;
+  }
+
+  synchronized (mutex) {
+    if (status != DRIVER_RUNNING) {
+      return status;
+    }
+
+    CHECK(process != nullptr);
+
+    dispatch(process, &SchedulerProcess::suppressOffers, roles);
 
     return status;
   }