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

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

commit 0166d4bac067b4deaf1c09a7ef72f017809b56a5
Author: Andrei Sekretenko <[email protected]>
AuthorDate: Fri Oct 4 18:18:14 2019 -0400

    Introduced a role tree helper to modify a role and all its ancestors.
    
    Review: https://reviews.apache.org/r/71498/
---
 src/master/allocator/mesos/hierarchical.cpp | 55 ++++++++++-------------------
 src/master/allocator/mesos/hierarchical.hpp |  8 +++++
 src/tests/consumption_metrics_tests.cpp     |  6 ++--
 3 files changed, 31 insertions(+), 38 deletions(-)

diff --git a/src/master/allocator/mesos/hierarchical.cpp 
b/src/master/allocator/mesos/hierarchical.cpp
index 0854956..467908c 100644
--- a/src/master/allocator/mesos/hierarchical.cpp
+++ b/src/master/allocator/mesos/hierarchical.cpp
@@ -354,16 +354,13 @@ void RoleTree::trackReservations(const Resources& 
resources)
     CHECK(Resources::isReserved(r));
 
     const string& reservationRole = Resources::reservationRole(r);
-
-    Role* current = &(*this)[reservationRole];
     ResourceQuantities quantities = ResourceQuantities::fromScalarResources(r);
 
-    // Track it hierarchically up to the root.
-    // Create new role tree node if necessary.
-    for (; current != nullptr; current = current->parent) {
+    // NOTE: If necessary, a new role tree node is created.
+    applyToRoleAndAncestors(&(*this)[reservationRole], [&](Role* current) {
       current->reservationScalarQuantities_ += quantities;
       updateQuotaConsumedMetric(current);
-    }
+    });
   }
 }
 
@@ -374,17 +371,14 @@ void RoleTree::untrackReservations(const Resources& 
resources)
     CHECK(Resources::isReserved(r));
 
     const string& reservationRole = Resources::reservationRole(r);
-    CHECK_CONTAINS(roles_, reservationRole);
-
     ResourceQuantities quantities = ResourceQuantities::fromScalarResources(r);
 
-    // Track it hierarchically up to the root.
-    for (Role* current = &(roles_.at(reservationRole)); current != nullptr;
-         current = current->parent) {
-      CHECK_CONTAINS(current->reservationScalarQuantities_, quantities);
-      current->reservationScalarQuantities_ -= quantities;
-      updateQuotaConsumedMetric(current);
-    }
+    applyToRoleAndAncestors(
+      CHECK_NOTNONE(get_(reservationRole)), [&](Role* current) {
+        CHECK_CONTAINS(current->reservationScalarQuantities_, quantities);
+        current->reservationScalarQuantities_ -= quantities;
+        updateQuotaConsumedMetric(current);
+      });
 
     tryRemove(reservationRole);
   }
@@ -397,14 +391,10 @@ void RoleTree::trackAllocated(const Resources& resources_)
       const string& role,
       const Resources& resources,
       resources_.scalars().allocations()) {
-    CHECK_CONTAINS(roles_, role);
-
-    // Track it hierarchically up to the root.
-    for (Role* current = &(roles_.at(role)); current != nullptr;
-         current = current->parent) {
+    applyToRoleAndAncestors(CHECK_NOTNONE(get_(role)), [&](Role* current) {
       current->allocatedScalars_ += resources;
       updateQuotaConsumedMetric(current);
-    }
+    });
   }
 }
 
@@ -415,15 +405,10 @@ void RoleTree::untrackAllocated(const Resources& 
resources_)
       const string& role,
       const Resources& resources,
       resources_.scalars().allocations()) {
-    CHECK_CONTAINS(roles_, role);
-
-    // Track it hierarchically up to the root.
-    for (Role* current = &(roles_.at(role)); current != nullptr;
-         current = current->parent) {
-      CHECK_CONTAINS(current->allocatedScalars_, resources);
+    applyToRoleAndAncestors(CHECK_NOTNONE(get_(role)), [&](Role* current) {
       current->allocatedScalars_ -= resources;
       updateQuotaConsumedMetric(current);
-    }
+    });
   }
 }
 
@@ -478,11 +463,10 @@ void RoleTree::trackOfferedOrAllocated(const Resources& 
resources_)
       const string& role,
       const Resources& resources,
       resources_.scalars().allocations()) {
-    // Track it hierarchically up to the root.
-    for (Role* current = CHECK_NOTNONE(get_(role)); current != nullptr;
-         current = current->parent) {
+    applyToRoleAndAncestors(
+        CHECK_NOTNONE(get_(role)), [&resources](Role* current) {
       current->offeredOrAllocatedScalars_ += resources;
-    }
+    });
   }
 }
 
@@ -497,14 +481,13 @@ void RoleTree::untrackOfferedOrAllocated(const Resources& 
resources_)
       const string& role,
       const Resources& resources,
       resources_.scalars().allocations()) {
-    // Untrack it hierarchically up to the root.
-    for (Role* current = CHECK_NOTNONE(get_(role)); current != nullptr;
-         current = current->parent) {
+    applyToRoleAndAncestors(
+        CHECK_NOTNONE(get_(role)), [&resources](Role* current) {
       CHECK_CONTAINS(current->offeredOrAllocatedScalars_, resources)
         << " Role: " << current->role
         << " offeredOrAllocated: " << current->offeredOrAllocatedScalars_;
       current->offeredOrAllocatedScalars_ -= resources;
-    }
+    });
   }
 }
 
diff --git a/src/master/allocator/mesos/hierarchical.hpp 
b/src/master/allocator/mesos/hierarchical.hpp
index 647718d..2648f66 100644
--- a/src/master/allocator/mesos/hierarchical.hpp
+++ b/src/master/allocator/mesos/hierarchical.hpp
@@ -267,6 +267,14 @@ private:
   // along the tree path will be created if necessary.
   Role& operator[](const std::string& role);
 
+  // Helper for modifying a role and all its ancestors.
+  template<class UnaryFunction>
+  static void applyToRoleAndAncestors(Role* role, UnaryFunction f) {
+    for (; role != nullptr; role = role->parent) {
+      f(role);
+    }
+  }
+
   // Try to remove the role associated with the given role.
   // The role must exist. The role and its ancestors will be removed
   // if they become "empty". See "Role:isEmpty()".
diff --git a/src/tests/consumption_metrics_tests.cpp 
b/src/tests/consumption_metrics_tests.cpp
index 6516aa2..d075bb6 100644
--- a/src/tests/consumption_metrics_tests.cpp
+++ b/src/tests/consumption_metrics_tests.cpp
@@ -105,7 +105,8 @@ TEST_F(ConsumptionMetricsTest, Launch)
 
   Future<Event::Offers> offers;
   EXPECT_CALL(*scheduler, offers(_, _))
-    .WillOnce(FutureArg<1>(&offers));
+    .WillOnce(FutureArg<1>(&offers))
+    .WillRepeatedly(Return()); // Ignore subsequent offers.
 
   // Subscribe the framework.
   v1::scheduler::TestMesos mesos(
@@ -212,7 +213,8 @@ TEST_F(ConsumptionMetricsTest, ReservedResource)
 
   Future<Event::Offers> offers;
   EXPECT_CALL(*scheduler, offers(_, _))
-    .WillOnce(FutureArg<1>(&offers));
+    .WillOnce(FutureArg<1>(&offers))
+    .WillRepeatedly(Return()); // Ignore subsequent offers.
 
   // Subscribe the framework.
   v1::scheduler::TestMesos mesos(

Reply via email to