You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mesos.apache.org by mz...@apache.org on 2019/07/01 21:20:42 UTC

[mesos] 04/04: Updated registry operation `UpdateQuota` to persist `QuotaConfig`.

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 3720e4cf5f7cb0d8e98afacea39528bd41c767b4
Author: Meng Zhu <mz...@mesosphere.io>
AuthorDate: Fri Jun 28 14:16:00 2019 -0700

    Updated registry operation `UpdateQuota` to persist `QuotaConfig`.
    
    The new operations will mutate the `quota_configs` field in
    the registry to persist `QuotaConfigs` configured by the new
    `UPDATE_QUOTA` call as well as the legacy `SET_QUOTA` and
    `REMOVE_QUOTA` calls.
    
    The operation removes any entries in the legacy `quotas` field
    with the same role name. In addition, it also adds/removes the
    minimum capability `QUOTA_V2` accordingly: if `quota_configs`
    is empty the capability will be removed otherwise it will
    be added.
    
    This operation replaces the `REMOVE_QUOTA` operation.
    
    Also fixed/disabled affected tests.
    
    Review: https://reviews.apache.org/r/70951
---
 src/master/master.cpp            |   8 +
 src/master/quota.cpp             |  91 ++++++----
 src/master/quota.hpp             |  35 +---
 src/master/quota_handler.cpp     |  43 ++++-
 src/tests/master_quota_tests.cpp |   4 +-
 src/tests/registrar_tests.cpp    | 381 ++++++++++++++++++++++++---------------
 6 files changed, 348 insertions(+), 214 deletions(-)

diff --git a/src/master/master.cpp b/src/master/master.cpp
index 826362d..5247377 100644
--- a/src/master/master.cpp
+++ b/src/master/master.cpp
@@ -1729,10 +1729,18 @@ Future<Nothing> Master::_recover(const Registry& registry)
   }
 
   // Save the quotas for each role.
+
+  // First recover from the legacy quota entries.
   foreach (const Registry::Quota& quota, registry.quotas()) {
     quotas[quota.info().role()] = Quota{quota.info()};
   }
 
+  // Then the new ones.
+  foreach (const quota::QuotaConfig& config, registry.quota_configs()) {
+    CHECK_NOT_CONTAINS(quotas, config.role());
+    quotas[config.role()] = Quota{config};
+  }
+
   // We notify the allocator via the `recover()` call. This has to be
   // done before the first agent reregisters and makes its resources
   // available for allocation. This is necessary because at this point
diff --git a/src/master/quota.cpp b/src/master/quota.cpp
index a7ee845..c9f39f3 100644
--- a/src/master/quota.cpp
+++ b/src/master/quota.cpp
@@ -29,9 +29,12 @@
 #include <stout/option.hpp>
 #include <stout/set.hpp>
 
+#include "common/protobuf_utils.hpp"
 #include "common/resources_utils.hpp"
 #include "common/validation.hpp"
 
+#include "master/constants.hpp"
+
 using google::protobuf::Map;
 using google::protobuf::RepeatedPtrField;
 
@@ -46,49 +49,75 @@ namespace internal {
 namespace master {
 namespace quota {
 
-UpdateQuota::UpdateQuota(const QuotaInfo& quotaInfo)
-  : info(quotaInfo) {}
+UpdateQuota::UpdateQuota(
+    const google::protobuf::RepeatedPtrField<QuotaConfig>& quotaConfigs)
+  : configs(quotaConfigs) {}
 
 
 Try<bool> UpdateQuota::perform(
-    Registry* registry,
-    hashset<SlaveID>* /*slaveIDs*/)
+    Registry* registry, hashset<SlaveID>* /*slaveIDs*/)
 {
-  // If there is already quota stored for the role, update the entry.
-  foreach (Registry::Quota& quota, *registry->mutable_quotas()) {
-    if (quota.info().role() == info.role()) {
-      quota.mutable_info()->CopyFrom(info);
-      return true; // Mutation.
+  google::protobuf::RepeatedPtrField<QuotaConfig>& registryConfigs =
+    *registry->mutable_quota_configs();
+
+  foreach (const QuotaConfig& config, configs) {
+    // Check if there is already quota stored for the role.
+    int configIndex = find_if(
+        registryConfigs.begin(),
+        registryConfigs.end(),
+        [&](const QuotaConfig& registryConfig) {
+          return registryConfig.role() == config.role();
+        }) -
+        registryConfigs.begin();
+
+    if (Quota(config) == DEFAULT_QUOTA) {
+      // Erase if present, otherwise no-op.
+      if (configIndex < registryConfigs.size()) {
+        registryConfigs.DeleteSubrange(configIndex, 1);
+      }
+    } else {
+      // Modify if present, otherwise insert.
+      if (configIndex < registryConfigs.size()) {
+        // TODO(mzhu): Check if we are setting quota to the same value.
+        // If so, no need to mutate.
+        *registryConfigs.Mutable(configIndex) = config;
+      } else {
+        *registryConfigs.Add() = config;
+      }
     }
-  }
-
-  // If there is no quota yet for the role, create a new entry.
-  registry->add_quotas()->mutable_info()->CopyFrom(info);
-
-  return true; // Mutation.
-}
 
+    // Remove the old `QuotaInfo` entries if any.
 
-RemoveQuota::RemoveQuota(const string& _role) : role(_role) {}
+    google::protobuf::RepeatedPtrField<Registry::Quota>& quotas =
+      *registry->mutable_quotas();
 
+    int quotaIndex = find_if(
+        quotas.begin(),
+        quotas.end(),
+        [&](const Registry::Quota& quota) {
+        return quota.info().role() == config.role();
+        }) -
+        quotas.begin();
 
-Try<bool> RemoveQuota::perform(
-    Registry* registry,
-    hashset<SlaveID>* /*slaveIDs*/)
-{
-  // Remove quota for the role if a corresponding entry exists.
-  for (int i = 0; i < registry->quotas().size(); ++i) {
-    const Registry::Quota& quota = registry->quotas(i);
-
-    if (quota.info().role() == role) {
-      registry->mutable_quotas()->DeleteSubrange(i, 1);
-
-      // NOTE: Multiple entries per role are not allowed.
-      return true; // Mutation.
+    if (quotaIndex < quotas.size()) {
+      quotas.DeleteSubrange(quotaIndex, 1);
     }
   }
 
-  return false;
+  // Update the minimum capability.
+  if (!registryConfigs.empty()) {
+    protobuf::master::addMinimumCapability(
+        registry->mutable_minimum_capabilities(),
+        MasterInfo::Capability::QUOTA_V2);
+  } else {
+    protobuf::master::removeMinimumCapability(
+        registry->mutable_minimum_capabilities(),
+        MasterInfo::Capability::QUOTA_V2);
+  }
+
+  // We always return true here since there is currently
+  // no optimization for no mutation case.
+  return true;
 }
 
 
diff --git a/src/master/quota.hpp b/src/master/quota.hpp
index 959e312..c828798 100644
--- a/src/master/quota.hpp
+++ b/src/master/quota.hpp
@@ -50,43 +50,22 @@ namespace quota {
  * Sets quota for a role. No assumptions are made here: the role may
  * be unknown to the master, or quota can be already set for the role.
  * If there is no quota stored for the role, a new entry is created,
- * otherwise an existing one is updated. This operation always mutates
- * the registry.
- *
- * TODO(alexr): Introduce equality operator in `Registry::Quota` or
- * `QuotaInfo` to avoid mutation in case of update to an equal value.
- * However, even if we return `false` (i.e. no mutation), the current
- * implementation of the registrar will still save the object again.
+ * otherwise an existing one is updated.
  */
 class UpdateQuota : public RegistryOperation
 {
 public:
-  explicit UpdateQuota(const mesos::quota::QuotaInfo& quotaInfo);
-
-protected:
-  Try<bool> perform(Registry* registry, hashset<SlaveID>* slaveIDs) override;
-
-private:
-  const mesos::quota::QuotaInfo info;
-};
-
-
-/**
- * Removes quota for a role. If there is no quota stored for the role,
- * no action is performed.
- *
- * TODO(alexr): Consider uniting this operation with `UpdateQuota`.
- */
-class RemoveQuota : public RegistryOperation
-{
-public:
-  explicit RemoveQuota(const std::string& _role);
+  // This operation needs to take in a `RepeatedPtrField` of `QuotaConfig`
+  // to ensure all-or-nothing quota updates for multiple roles.
+  explicit UpdateQuota(
+      const google::protobuf::RepeatedPtrField<mesos::quota::QuotaConfig>&
+        quotaConfigs);
 
 protected:
   Try<bool> perform(Registry* registry, hashset<SlaveID>* slaveIDs) override;
 
 private:
-  const std::string role;
+  google::protobuf::RepeatedPtrField<mesos::quota::QuotaConfig> configs;
 };
 
 
diff --git a/src/master/quota_handler.cpp b/src/master/quota_handler.cpp
index 476e87e..f9c7431 100644
--- a/src/master/quota_handler.cpp
+++ b/src/master/quota_handler.cpp
@@ -48,6 +48,8 @@
 
 namespace http = process::http;
 
+using google::protobuf::RepeatedPtrField;
+
 using http::Accepted;
 using http::BadRequest;
 using http::Conflict;
@@ -657,9 +659,29 @@ Future<http::Response> Master::QuotaHandler::__set(
   // fails because in this case the master fails as well.
   master->quotas[quotaInfo.role()] = quota;
 
+  // Construct `RepeatedPtrField<QuotaConfig>` from the legacy `QuotaInfo`
+  // for forward compatibility.
+  RepeatedPtrField<QuotaConfig> configs = [&quotaInfo]() {
+    QuotaConfig config;
+    *config.mutable_role() = quotaInfo.role();
+
+    google::protobuf::Map<string, Value::Scalar> quota;
+    foreach (const Resource& r, quotaInfo.guarantee()) {
+      quota[r.name()] = r.scalar();
+    }
+
+    *config.mutable_guarantees() = quota;
+    *config.mutable_limits() = std::move(quota);
+
+    RepeatedPtrField<QuotaConfig> configs;
+    *configs.Add() = std::move(config);
+
+    return configs;
+  }();
+
   // Update the registry with the new quota and acknowledge the request.
-  return master->registrar->apply(Owned<RegistryOperation>(
-      new quota::UpdateQuota(quotaInfo)))
+  return master->registrar
+    ->apply(Owned<RegistryOperation>(new quota::UpdateQuota(configs)))
     .then(defer(master->self(), [=](bool result) -> Future<http::Response> {
       // See the top comment in "master/quota.hpp" for why this check is here.
       CHECK(result);
@@ -797,9 +819,22 @@ Future<http::Response> Master::QuotaHandler::__remove(const string& role) const
   // will be restored automatically during the recovery.
   master->quotas.erase(role);
 
+  // Remove quota is equivalent to configure quota back to the default.
+  // We need to wrap it up in `RepeatedPtrField<QuotaConfig>` for
+  // foward compatibility.
+  RepeatedPtrField<QuotaConfig> configs = [&role]() {
+    QuotaConfig config;
+    *config.mutable_role() = role;
+
+    RepeatedPtrField<QuotaConfig> configs;
+    *configs.Add() = std::move(config);
+
+    return configs;
+  }();
+
   // Update the registry with the removed quota and acknowledge the request.
-  return master->registrar->apply(Owned<RegistryOperation>(
-      new quota::RemoveQuota(role)))
+  return master->registrar
+    ->apply(Owned<RegistryOperation>(new quota::UpdateQuota(configs)))
     .then(defer(master->self(), [=](bool result) -> Future<http::Response> {
       // See the top comment in "master/quota.hpp" for why this check is here.
       CHECK(result);
diff --git a/src/tests/master_quota_tests.cpp b/src/tests/master_quota_tests.cpp
index 4b805e9..34a6520 100644
--- a/src/tests/master_quota_tests.cpp
+++ b/src/tests/master_quota_tests.cpp
@@ -1268,7 +1268,9 @@ TEST_F(MasterQuotaTest, AvailableResourcesAfterRescinding)
 
 // Checks that quota is recovered correctly after master failover if
 // the expected size of the cluster is zero.
-TEST_F(MasterQuotaTest, RecoverQuotaEmptyCluster)
+//
+// TODO(mzhu): Enable this once master is QUOTA_V2 capable.
+TEST_F(MasterQuotaTest, DISABLED_RecoverQuotaEmptyCluster)
 {
   master::Flags masterFlags = CreateMasterFlags();
   masterFlags.registry = "replicated_log";
diff --git a/src/tests/registrar_tests.cpp b/src/tests/registrar_tests.cpp
index 81979d7..1343b5e 100644
--- a/src/tests/registrar_tests.cpp
+++ b/src/tests/registrar_tests.cpp
@@ -21,6 +21,8 @@
 #include <string>
 #include <vector>
 
+#include <google/protobuf/util/message_differencer.h>
+
 #include <mesos/attributes.hpp>
 #include <mesos/type_utils.hpp>
 
@@ -83,6 +85,7 @@ using process::http::Response;
 using process::http::Unauthorized;
 
 using google::protobuf::RepeatedPtrField;
+using google::protobuf::util::MessageDifferencer;
 
 using mesos::internal::protobuf::maintenance::createMachineList;
 using mesos::internal::protobuf::maintenance::createSchedule;
@@ -911,39 +914,50 @@ TEST_F(RegistrarTest, StopMaintenance)
 // Tests that adding and updating quotas in the registry works properly.
 TEST_F(RegistrarTest, UpdateQuota)
 {
-  const string ROLE1 = "role1";
-  const string ROLE2 = "role2";
-
-  // NOTE: `quotaResources1` yields a collection with two `Resource`
-  // objects once converted to `RepeatedPtrField`.
-  Resources quotaResources1 = Resources::parse("cpus:1;mem:1024").get();
-  Resources quotaResources2 = Resources::parse("cpus:2").get();
-
-  // Prepare `QuotaInfo` protobufs used in the test.
-  QuotaInfo quota1;
-  quota1.set_role(ROLE1);
-  quota1.mutable_guarantee()->CopyFrom(quotaResources1);
-
-  Option<Error> validateError1 = quota::validation::quotaInfo(quota1);
-  EXPECT_NONE(validateError1);
-
-  QuotaInfo quota2;
-  quota2.set_role(ROLE2);
-  quota2.mutable_guarantee()->CopyFrom(quotaResources1);
-
-  Option<Error> validateError2 = quota::validation::quotaInfo(quota2);
-  EXPECT_NONE(validateError2);
+  // Helper to construct `QuotaConfig`.
+  auto createQuotaConfig = [](const string& role,
+                              const string& quantitiesString,
+                              const string& limitsString) {
+    QuotaConfig config;
+    config.set_role(role);
+
+    google::protobuf::Map<string, Value::Scalar> guarantees_;
+    ResourceQuantities quantities =
+      CHECK_NOTERROR(ResourceQuantities::fromString(quantitiesString));
+    foreachpair (const string& name, const Value::Scalar& scalar, quantities) {
+      guarantees_[name] = scalar;
+    }
+
+    google::protobuf::Map<string, Value::Scalar> limits_;
+    ResourceLimits limits =
+      CHECK_NOTERROR(ResourceLimits::fromString(limitsString));
+    foreachpair (const string& name, const Value::Scalar& scalar, limits) {
+      limits_[name] = scalar;
+    }
+
+    *config.mutable_guarantees() = std::move(guarantees_);
+    *config.mutable_limits() = std::move(limits_);
+
+    return config;
+  };
+
+  RepeatedPtrField<QuotaConfig> configs;
 
   {
-    // Prepare the registrar; see the comment above why we need to do this in
-    // every scope.
+    // Initially no quota and minimum capabilities are recorded.
+
     Registrar registrar(flags, state);
     Future<Registry> registry = registrar.recover(master);
     AWAIT_READY(registry);
 
-    // Store quota for a role without quota.
-    AWAIT_TRUE(registrar.apply(Owned<RegistryOperation>(
-        new UpdateQuota(quota1))));
+    EXPECT_EQ(0, registry->quota_configs().size());
+    EXPECT_EQ(0, registry->minimum_capabilities().size());
+
+    // Store quota for a role with default quota.
+    *configs.Add() = createQuotaConfig("role1", "", "");
+
+    AWAIT_TRUE(
+        registrar.apply(Owned<RegistryOperation>(new UpdateQuota(configs))));
   }
 
   {
@@ -951,20 +965,17 @@ TEST_F(RegistrarTest, UpdateQuota)
     Future<Registry> registry = registrar.recover(master);
     AWAIT_READY(registry);
 
-    // Check that the recovered quota matches the one we stored.
-    ASSERT_EQ(1, registry->quotas().size());
-    EXPECT_EQ(ROLE1, registry->quotas(0).info().role());
-    ASSERT_EQ(2, registry->quotas(0).info().guarantee().size());
+    // Default quota is not persisted into the registry.
+    EXPECT_EQ(0, registry->quota_configs().size());
+    EXPECT_EQ(0, registry->minimum_capabilities().size());
 
-    Resources storedResources(registry->quotas(0).info().guarantee());
-    EXPECT_EQ(quotaResources1, storedResources);
+    // Update quota for `role1`.
+    configs.Clear();
+    *configs.Add() =
+      createQuotaConfig("role1", "cpus:1;mem:1024", "cpus:2;mem:2048");
 
-    // Change quota for `ROLE1`.
-    quota1.mutable_guarantee()->CopyFrom(quotaResources2);
-
-    // Update the only stored quota.
-    AWAIT_TRUE(registrar.apply(Owned<RegistryOperation>(
-        new UpdateQuota(quota1))));
+    AWAIT_TRUE(
+        registrar.apply(Owned<RegistryOperation>(new UpdateQuota(configs))));
   }
 
   {
@@ -972,17 +983,48 @@ TEST_F(RegistrarTest, UpdateQuota)
     Future<Registry> registry = registrar.recover(master);
     AWAIT_READY(registry);
 
-    // Check that the recovered quota matches the one we updated.
-    ASSERT_EQ(1, registry->quotas().size());
-    EXPECT_EQ(ROLE1, registry->quotas(0).info().role());
-    ASSERT_EQ(1, registry->quotas(0).info().guarantee().size());
-
-    Resources storedResources(registry->quotas(0).info().guarantee());
-    EXPECT_EQ(quotaResources2, storedResources);
-
-    // Store one more quota for a role without quota.
-    AWAIT_TRUE(registrar.apply(Owned<RegistryOperation>(
-        new UpdateQuota(quota2))));
+    EXPECT_EQ(1, registry->quota_configs().size());
+
+    JSON::Array expected = CHECK_NOTERROR(JSON::parse<JSON::Array>(
+    R"~(
+    [
+      {
+        "guarantees": {
+          "cpus": {
+            "value": 1
+          },
+          "mem": {
+            "value": 1024
+          }
+        },
+        "limits": {
+          "cpus": {
+            "value": 2
+          },
+          "mem": {
+            "value": 2048
+          }
+        },
+        "role": "role1"
+      }
+    ])~"));
+
+    EXPECT_EQ(expected, JSON::protobuf(registry->quota_configs()));
+
+    // The `QUOTA_V2` capability is added to the registry.
+    //
+    // TODO(mzhu): This assumes the the registry starts empty which might not
+    // be in the future. Just check the presence of `QUOTA_V2`.
+    EXPECT_EQ(1, registry->minimum_capabilities().size());
+    EXPECT_EQ("QUOTA_V2", registry->minimum_capabilities(0).capability());
+
+    // Update quota for "role2".
+    configs.Clear();
+    *configs.Add() =
+      createQuotaConfig("role2", "cpus:1;mem:1024", "cpus:2;mem:2048");
+
+    AWAIT_TRUE(
+        registrar.apply(Owned<RegistryOperation>(new UpdateQuota(configs))));
   }
 
   {
@@ -994,88 +1036,63 @@ TEST_F(RegistrarTest, UpdateQuota)
     // NOTE: We assume quota messages are stored in order they have
     // been added.
     // TODO(alexr): Consider removing dependency on the order.
-    ASSERT_EQ(2, registry->quotas().size());
-    EXPECT_EQ(ROLE1, registry->quotas(0).info().role());
-    ASSERT_EQ(1, registry->quotas(0).info().guarantee().size());
-
-    EXPECT_EQ(ROLE2, registry->quotas(1).info().role());
-    ASSERT_EQ(2, registry->quotas(1).info().guarantee().size());
-
-    Resources storedResources(registry->quotas(1).info().guarantee());
-    EXPECT_EQ(quotaResources1, storedResources);
-
-    // Change quota for `role2`.
-    quota2.mutable_guarantee()->CopyFrom(quotaResources2);
-
-    // Update quota for `role2` in presence of multiple quotas.
-    AWAIT_TRUE(registrar.apply(Owned<RegistryOperation>(
-        new UpdateQuota(quota2))));
-  }
-
-  {
-    Registrar registrar(flags, state);
-    Future<Registry> registry = registrar.recover(master);
-    AWAIT_READY(registry);
-
-    // Check that the recovered quotas match those we stored and updated
-    // previously.
-    // NOTE: We assume quota messages are stored in order they have been
-    // added and update does not change the order.
-    // TODO(alexr): Consider removing dependency on the order.
-    ASSERT_EQ(2, registry->quotas().size());
-
-    EXPECT_EQ(ROLE1, registry->quotas(0).info().role());
-    ASSERT_EQ(1, registry->quotas(0).info().guarantee().size());
-
-    Resources storedResources1(registry->quotas(0).info().guarantee());
-    EXPECT_EQ(quotaResources2, storedResources1);
-
-    EXPECT_EQ(ROLE2, registry->quotas(1).info().role());
-    ASSERT_EQ(1, registry->quotas(1).info().guarantee().size());
-
-    Resources storedResources2(registry->quotas(1).info().guarantee());
-    EXPECT_EQ(quotaResources2, storedResources2);
-  }
-}
-
-
-// Tests removing quotas from the registry.
-TEST_F(RegistrarTest, RemoveQuota)
-{
-  const string ROLE1 = "role1";
-  const string ROLE2 = "role2";
-
-  {
-    // Prepare the registrar; see the comment above why we need to do this in
-    // every scope.
-    Registrar registrar(flags, state);
-    Future<Registry> registry = registrar.recover(master);
-    AWAIT_READY(registry);
-
-    // NOTE: `quotaResources` yields a collection with two `Resource`
-    // objects once converted to `RepeatedPtrField`.
-    Resources quotaResources1 = Resources::parse("cpus:1;mem:1024").get();
-    Resources quotaResources2 = Resources::parse("cpus:2").get();
-
-    // Prepare `QuotaInfo` protobufs.
-    QuotaInfo quota1;
-    quota1.set_role(ROLE1);
-    quota1.mutable_guarantee()->CopyFrom(quotaResources1);
-
-    Option<Error> validateError1 = quota::validation::quotaInfo(quota1);
-    EXPECT_NONE(validateError1);
-
-    QuotaInfo quota2;
-    quota2.set_role(ROLE2);
-    quota2.mutable_guarantee()->CopyFrom(quotaResources2);
-
-    Option<Error> validateError2 = quota::validation::quotaInfo(quota2);
-    EXPECT_NONE(validateError2);
-
-    AWAIT_TRUE(registrar.apply(Owned<RegistryOperation>(
-        new UpdateQuota(quota1))));
-    AWAIT_TRUE(registrar.apply(Owned<RegistryOperation>(
-        new UpdateQuota(quota2))));
+    JSON::Array expected = CHECK_NOTERROR(JSON::parse<JSON::Array>(
+    R"~(
+    [
+      {
+        "guarantees": {
+          "cpus": {
+            "value": 1
+          },
+          "mem": {
+            "value": 1024
+          }
+        },
+        "limits": {
+          "cpus": {
+            "value": 2
+          },
+          "mem": {
+            "value": 2048
+          }
+        },
+        "role": "role1"
+      },
+      {
+        "guarantees": {
+          "cpus": {
+            "value": 1
+          },
+          "mem": {
+            "value": 1024
+          }
+        },
+        "limits": {
+          "cpus": {
+            "value": 2
+          },
+          "mem": {
+            "value": 2048
+          }
+        },
+        "role": "role2"
+      }
+    ])~"));
+
+    EXPECT_EQ(expected, JSON::protobuf(registry->quota_configs()));
+
+    // TODO(mzhu): This assumes the the registry starts empty which might not
+    // be in the future. Just check the presence of `QUOTA_V2`.
+    EXPECT_EQ(1, registry->minimum_capabilities().size());
+    EXPECT_EQ("QUOTA_V2", registry->minimum_capabilities(0).capability());
+
+    // Change quota for "role1"` and "role2"` in a single call.
+    configs.Clear();
+    *configs.Add() = createQuotaConfig("role1", "cpus:2", "cpus:4");
+    *configs.Add() = createQuotaConfig("role2", "cpus:2", "cpus:4");
+
+    AWAIT_TRUE(
+        registrar.apply(Owned<RegistryOperation>(new UpdateQuota(configs))));
   }
 
   {
@@ -1084,16 +1101,53 @@ TEST_F(RegistrarTest, RemoveQuota)
     AWAIT_READY(registry);
 
     // Check that the recovered quotas match those we stored previously.
-    // NOTE: We assume quota messages are stored in order they have been
-    // added.
+    // NOTE: We assume quota messages are stored in order they have
+    // been added.
     // TODO(alexr): Consider removing dependency on the order.
-    ASSERT_EQ(2, registry->quotas().size());
-    EXPECT_EQ(ROLE1, registry->quotas(0).info().role());
-    EXPECT_EQ(ROLE2, registry->quotas(1).info().role());
-
-    // Remove quota for `role2`.
-    AWAIT_TRUE(registrar.apply(Owned<RegistryOperation>(
-        new RemoveQuota(ROLE2))));
+    JSON::Array expected = CHECK_NOTERROR(JSON::parse<JSON::Array>(
+    R"~(
+    [
+      {
+        "guarantees": {
+          "cpus": {
+            "value": 2
+          }
+        },
+        "limits": {
+          "cpus": {
+            "value": 4
+          }
+        },
+        "role": "role1"
+      },
+      {
+        "guarantees": {
+          "cpus": {
+            "value": 2
+          }
+        },
+        "limits": {
+          "cpus": {
+            "value": 4
+          }
+        },
+        "role": "role2"
+      }
+    ])~"));
+
+    EXPECT_EQ(expected, JSON::protobuf(registry->quota_configs()));
+
+    // TODO(mzhu): This assumes the the registry starts empty which might not
+    // be in the future. Just check the presence of `QUOTA_V2`.
+    EXPECT_EQ(1, registry->minimum_capabilities().size());
+    EXPECT_EQ("QUOTA_V2", registry->minimum_capabilities(0).capability());
+
+    // Reset "role2"` quota to default.
+    configs.Clear();
+    *configs.Add() = createQuotaConfig("role2", "", "");
+
+    AWAIT_TRUE(
+        registrar.apply(Owned<RegistryOperation>(new UpdateQuota(configs))));
   }
 
   {
@@ -1101,13 +1155,39 @@ TEST_F(RegistrarTest, RemoveQuota)
     Future<Registry> registry = registrar.recover(master);
     AWAIT_READY(registry);
 
-    // Check that there is only one quota left in the registry.
-    ASSERT_EQ(1, registry->quotas().size());
-    EXPECT_EQ(ROLE1, registry->quotas(0).info().role());
-
-    // Remove quota for `ROLE1`.
-    AWAIT_TRUE(registrar.apply(Owned<RegistryOperation>(
-        new RemoveQuota(ROLE1))));
+    configs.Clear();
+    *configs.Add() = createQuotaConfig("role1", "cpus:2", "cpus:4");
+    JSON::Array expected = CHECK_NOTERROR(JSON::parse<JSON::Array>(
+    R"~(
+    [
+      {
+        "guarantees": {
+          "cpus": {
+            "value": 2
+          }
+        },
+        "limits": {
+          "cpus": {
+            "value": 4
+          }
+        },
+        "role": "role1"
+      }
+    ])~"));
+
+    EXPECT_EQ(expected, JSON::protobuf(registry->quota_configs()));
+
+    // TODO(mzhu): This assumes the the registry starts empty which might not
+    // be in the future. Just check the presence of `QUOTA_V2`.
+    EXPECT_EQ(1, registry->minimum_capabilities().size());
+    EXPECT_EQ("QUOTA_V2", registry->minimum_capabilities(0).capability());
+
+    // Reset "role1"` quota to default.
+    configs.Clear();
+    *configs.Add() = createQuotaConfig("role1", "", "");
+
+    AWAIT_TRUE(
+        registrar.apply(Owned<RegistryOperation>(new UpdateQuota(configs))));
   }
 
   {
@@ -1115,8 +1195,9 @@ TEST_F(RegistrarTest, RemoveQuota)
     Future<Registry> registry = registrar.recover(master);
     AWAIT_READY(registry);
 
-    // Check that there are no more quotas at this point.
-    ASSERT_TRUE(registry->quotas().empty());
+    EXPECT_EQ(0, registry->quota_configs().size());
+    // The `QUOTA_V2` capability is removed because `quota_configs` is empty.
+    EXPECT_EQ(0, registry->minimum_capabilities().size());
   }
 }