Maintenance Primitives: Added maintenance schedule endpoint.

Endpoint: /maintenance/schedule

Registry operation = maintenance::UpdateSchedule
  Replaces the schedule with the given one.  Also sets all scheduled
  machines into Draining mode.

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


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

Branch: refs/heads/master
Commit: 3b0fe5cc047f9b96dcd96771fa8945483ce92e72
Parents: 8e43aba
Author: Joseph Wu <[email protected]>
Authored: Sun Aug 30 13:56:02 2015 -0400
Committer: Joris Van Remoortere <[email protected]>
Committed: Mon Aug 31 13:09:48 2015 -0400

----------------------------------------------------------------------
 src/Makefile.am                        |   1 +
 src/master/http.cpp                    | 126 +++++++++++++++
 src/master/maintenance.cpp             |  76 +++++++++
 src/master/maintenance.hpp             |  31 ++++
 src/master/master.cpp                  |   6 +
 src/master/master.hpp                  |   5 +
 src/tests/master_maintenance_tests.cpp | 239 ++++++++++++++++++++++++++++
 src/tests/registrar_tests.cpp          | 184 +++++++++++++++++++++
 8 files changed, 668 insertions(+)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/mesos/blob/3b0fe5cc/src/Makefile.am
----------------------------------------------------------------------
diff --git a/src/Makefile.am b/src/Makefile.am
index 1b07af4..7b4d9f6 100644
--- a/src/Makefile.am
+++ b/src/Makefile.am
@@ -1639,6 +1639,7 @@ mesos_tests_SOURCES =                                     
        \
   tests/master_allocator_tests.cpp                             \
   tests/master_authorization_tests.cpp                         \
   tests/master_contender_detector_tests.cpp                    \
+  tests/master_maintenance_tests.cpp                           \
   tests/master_slave_reconciliation_tests.cpp                  \
   tests/master_tests.cpp                                       \
   tests/master_validation_tests.cpp                            \

http://git-wip-us.apache.org/repos/asf/mesos/blob/3b0fe5cc/src/master/http.cpp
----------------------------------------------------------------------
diff --git a/src/master/http.cpp b/src/master/http.cpp
index 37d76ee..44178d8 100644
--- a/src/master/http.cpp
+++ b/src/master/http.cpp
@@ -30,19 +30,28 @@
 
 #include <mesos/authorizer/authorizer.hpp>
 
+#include <mesos/maintenance/maintenance.hpp>
+
+#include <process/defer.hpp>
 #include <process/help.hpp>
 
 #include <process/metrics/metrics.hpp>
 
 #include <stout/base64.hpp>
 #include <stout/foreach.hpp>
+#include <stout/hashmap.hpp>
+#include <stout/hashset.hpp>
 #include <stout/json.hpp>
 #include <stout/lambda.hpp>
 #include <stout/net.hpp>
+#include <stout/nothing.hpp>
 #include <stout/numify.hpp>
 #include <stout/os.hpp>
+#include <stout/protobuf.hpp>
 #include <stout/result.hpp>
 #include <stout/strings.hpp>
+#include <stout/try.hpp>
+#include <stout/utils.hpp>
 
 #include "common/attributes.hpp"
 #include "common/build.hpp"
@@ -53,6 +62,7 @@
 
 #include "logging/logging.hpp"
 
+#include "master/maintenance.hpp"
 #include "master/master.hpp"
 #include "master/validation.hpp"
 
@@ -98,6 +108,7 @@ using mesos::internal::model;
 // Pull in definitions from process.
 using process::http::Response;
 using process::http::Request;
+using process::Owned;
 
 
 // TODO(bmahler): Kill these in favor of automatic Proto->JSON Conversion (when
@@ -1349,6 +1360,121 @@ Future<Response> Master::Http::tasks(const Request& 
request) const
 }
 
 
+// /master/maintenance/schedule endpoint help.
+const string Master::Http::MAINTENANCE_SCHEDULE_HELP = HELP(
+    TLDR(
+        "Returns or updates the cluster's maintenance schedule."),
+    USAGE(
+        "/master/maintenance/schedule"),
+    DESCRIPTION(
+        "GET: Returns the current maintenance schedule as JSON.",
+        "POST: Validates the request body as JSON",
+        "  and updates the maintenance schedule."));
+
+
+// /master/maintenance/schedule endpoint handler.
+Future<Response> Master::Http::maintenanceSchedule(const Request& request) 
const
+{
+  if (request.method != "GET" && request.method != "POST") {
+    return BadRequest("Expecting GET or POST, got '" + request.method + "'");
+  }
+
+  // JSON-ify and return the current maintenance schedule.
+  if (request.method == "GET") {
+    // TODO(josephw): Return more than one schedule.
+    const mesos::maintenance::Schedule schedule =
+      master->maintenance.schedules.empty() ?
+        mesos::maintenance::Schedule() :
+        master->maintenance.schedules.front();
+
+    return OK(JSON::Protobuf(schedule), request.query.get("jsonp"));
+  }
+
+  // Parse the POST body as JSON.
+  Try<JSON::Object> jsonSchedule = JSON::parse<JSON::Object>(request.body);
+  if (jsonSchedule.isError()) {
+    return BadRequest(jsonSchedule.error());
+  }
+
+  // Convert the schedule to a protobuf.
+  Try<mesos::maintenance::Schedule> protoSchedule =
+    ::protobuf::parse<mesos::maintenance::Schedule>(jsonSchedule.get());
+
+  if (protoSchedule.isError()) {
+    return BadRequest(protoSchedule.error());
+  }
+
+  // Validate that the schedule only transitions machines between
+  // `UP` and `DRAINING` modes.
+  mesos::maintenance::Schedule schedule = protoSchedule.get();
+  Try<Nothing> isValid = maintenance::validation::schedule(
+      schedule,
+      master->machineInfos);
+
+  if (isValid.isError()) {
+    return BadRequest(isValid.error());
+  }
+
+  return master->registrar->apply(Owned<Operation>(
+      new maintenance::UpdateSchedule(schedule)))
+    .then(defer(master->self(), [=](bool result) -> Future<Response> {
+      // See the top comment in "master/maintenance.hpp" for why this check
+      // is here, and is appropriate.
+      CHECK(result);
+
+      // Update the master's local state with the new schedule.
+      // NOTE: We only add or remove differences between the current schedule
+      // and the new schedule.  This is because the `MachineInfo` struct
+      // holds more information than a maintenance schedule.
+      // For example, the `mode` field is not part of a maintenance schedule.
+
+      // TODO(josephw): allow more than one schedule.
+
+      // Put the machines in the updated schedule into a set.
+      // Save the unavailability, to help with updating some machines.
+      hashmap<MachineID, Unavailability> updated;
+      foreach (const mesos::maintenance::Window& window, schedule.windows()) {
+        foreach (const MachineID& id, window.machine_ids()) {
+          updated[id] = window.unavailability();
+        }
+      }
+
+      // NOTE: Copies are needed because this loop modifies the container.
+      foreachkey (const MachineID& id, utils::copy(master->machineInfos)) {
+        // Update the entry for each updated machine.
+        if (updated.contains(id)) {
+          master->machineInfos[id]
+            .mutable_unavailability()->CopyFrom(updated[id]);
+
+          continue;
+        }
+
+        // Delete the entry for each removed machine.
+        master->machineInfos.erase(id);
+      }
+
+      // Save each new machine, with the unavailability
+      // and starting in `DRAINING` mode.
+      foreach (const mesos::maintenance::Window& window, schedule.windows()) {
+        foreach (const MachineID& id, window.machine_ids()) {
+          MachineInfo info;
+          info.mutable_id()->CopyFrom(id);
+          info.set_mode(MachineInfo::DRAINING);
+          info.mutable_unavailability()->CopyFrom(window.unavailability());
+
+          master->machineInfos[id] = info;
+        }
+      }
+
+      // Replace the old schedule(s) with the new schedule.
+      master->maintenance.schedules.clear();
+      master->maintenance.schedules.push_back(schedule);
+
+      return OK();
+    }));
+}
+
+
 Result<Credential> Master::Http::authenticate(const Request& request) const
 {
   // By default, assume everyone is authenticated if no credentials

http://git-wip-us.apache.org/repos/asf/mesos/blob/3b0fe5cc/src/master/maintenance.cpp
----------------------------------------------------------------------
diff --git a/src/master/maintenance.cpp b/src/master/maintenance.cpp
index 221dd02..798c026 100644
--- a/src/master/maintenance.cpp
+++ b/src/master/maintenance.cpp
@@ -22,6 +22,7 @@
 
 #include <stout/duration.hpp>
 #include <stout/error.hpp>
+#include <stout/hashmap.hpp>
 #include <stout/hashset.hpp>
 #include <stout/ip.hpp>
 #include <stout/nothing.hpp>
@@ -36,6 +37,81 @@ namespace maintenance {
 
 using namespace mesos::maintenance;
 
+UpdateSchedule::UpdateSchedule(
+    const maintenance::Schedule& _schedule)
+  : schedule(_schedule) {}
+
+
+Try<bool> UpdateSchedule::perform(
+    Registry* registry,
+    hashset<SlaveID>* slaveIDs,
+    bool strict)
+{
+  // Put the machines in the existing schedule into a set.
+  hashset<MachineID> existing;
+  foreach (const maintenance::Schedule& agenda, registry->schedules()) {
+    foreach (const maintenance::Window& window, agenda.windows()) {
+      foreach (const MachineID& id, window.machine_ids()) {
+        existing.insert(id);
+      }
+    }
+  }
+
+  // Put the machines in the updated schedule into a set.
+  // Keep the relevant unavailability to help update existing machines.
+  hashmap<MachineID, Unavailability> updated;
+  foreach (const maintenance::Window& window, schedule.windows()) {
+    foreach (const MachineID& id, window.machine_ids()) {
+      updated[id] = window.unavailability();
+    }
+  }
+
+  // This operation overwrites the existing schedules with a new one.
+  // At the same time, every machine in the schedule is saved as a
+  // `MachineInfo` in the registry.
+
+  // TODO(josephw): allow more than one schedule.
+
+  // Loop through and modify the existing `MachineInfo` entries.
+  for (int i = registry->machines().machines().size() - 1; i >= 0; i--) {
+    const MachineID& id = registry->machines().machines(i).info().id();
+
+    // Update the `MachineInfo` entry for all machines in the schedule.
+    if (updated.contains(id)) {
+      registry->mutable_machines()->mutable_machines(i)->mutable_info()
+        ->mutable_unavailability()->CopyFrom(updated[id]);
+
+      continue;
+    }
+
+    // Delete the `MachineInfo` entry for each removed machine.
+    registry->mutable_machines()->mutable_machines()->DeleteSubrange(i, 1);
+  }
+
+  // Create new `MachineInfo` entries for each new machine.
+  foreach (const maintenance::Window& window, schedule.windows()) {
+    foreach (const MachineID& id, window.machine_ids()) {
+      if (existing.contains(id)) {
+        continue;
+      }
+
+      // Each newly scheduled machine starts in `DRAINING` mode.
+      Registry::Machine* machine = 
registry->mutable_machines()->add_machines();
+      MachineInfo* info = machine->mutable_info();
+      info->mutable_id()->CopyFrom(id);
+      info->set_mode(MachineInfo::DRAINING);
+      info->mutable_unavailability()->CopyFrom(window.unavailability());
+    }
+  }
+
+  // Replace the old schedule with the new schedule.
+  registry->clear_schedules();
+  registry->add_schedules()->CopyFrom(schedule);
+
+  return true; // Mutation.
+}
+
+
 namespace validation {
 
 Try<Nothing> schedule(

http://git-wip-us.apache.org/repos/asf/mesos/blob/3b0fe5cc/src/master/maintenance.hpp
----------------------------------------------------------------------
diff --git a/src/master/maintenance.hpp b/src/master/maintenance.hpp
index 833665e..42b5f9e 100644
--- a/src/master/maintenance.hpp
+++ b/src/master/maintenance.hpp
@@ -34,6 +34,37 @@ namespace mesos {
 namespace internal {
 namespace master {
 namespace maintenance {
+
+// Maintenance registry operations will never report a failure while
+// performing the operation (i.e. `perform` never returns an `Error`).
+// This means the underlying `Future` for the operation will always be
+// set to true.  If there is a separate failure, such as a network
+// partition, while performing the operation, this future will not be set.
+
+/**
+ * Updates the maintanence schedule of the cluster.  This transitions machines
+ * between `UP` and `DRAINING` modes only.  The given schedule must only
+ * add valid machines and remove machines that are not `DOWN`.
+ *
+ * TODO(josephw): allow more than one schedule.
+ */
+class UpdateSchedule : public Operation
+{
+public:
+  explicit UpdateSchedule(
+      const mesos::maintenance::Schedule& _schedule);
+
+protected:
+  Try<bool> perform(
+      Registry* registry,
+      hashset<SlaveID>* slaveIDs,
+      bool strict);
+
+private:
+  const mesos::maintenance::Schedule schedule;
+};
+
+
 namespace validation {
 
 /**

http://git-wip-us.apache.org/repos/asf/mesos/blob/3b0fe5cc/src/master/master.cpp
----------------------------------------------------------------------
diff --git a/src/master/master.cpp b/src/master/master.cpp
index 564fbcb..ea556f9 100644
--- a/src/master/master.cpp
+++ b/src/master/master.cpp
@@ -815,6 +815,12 @@ void Master::initialize()
           Http::log(request);
           return http.tasks(request);
         });
+  route("/maintenance/schedule",
+        Http::MAINTENANCE_SCHEDULE_HELP,
+        [http](const process::http::Request& request) {
+          Http::log(request);
+          return http.maintenanceSchedule(request);
+        });
 
   // Provide HTTP assets from a "webui" directory. This is either
   // specified via flags (which is necessary for running out of the

http://git-wip-us.apache.org/repos/asf/mesos/blob/3b0fe5cc/src/master/master.hpp
----------------------------------------------------------------------
diff --git a/src/master/master.hpp b/src/master/master.hpp
index edcbeed..175e623 100644
--- a/src/master/master.hpp
+++ b/src/master/master.hpp
@@ -832,6 +832,10 @@ private:
     process::Future<process::http::Response> tasks(
         const process::http::Request& request) const;
 
+    // /master/maintenance/schedule
+    process::Future<process::http::Response> maintenanceSchedule(
+        const process::http::Request& request) const;
+
     const static std::string SCHEDULER_HELP;
     const static std::string HEALTH_HELP;
     const static std::string OBSERVE_HELP;
@@ -842,6 +846,7 @@ private:
     const static std::string STATE_HELP;
     const static std::string STATESUMMARY_HELP;
     const static std::string TASKS_HELP;
+    const static std::string MAINTENANCE_SCHEDULE_HELP;
 
   private:
     // Helper for doing authentication, returns the credential used if

http://git-wip-us.apache.org/repos/asf/mesos/blob/3b0fe5cc/src/tests/master_maintenance_tests.cpp
----------------------------------------------------------------------
diff --git a/src/tests/master_maintenance_tests.cpp 
b/src/tests/master_maintenance_tests.cpp
new file mode 100644
index 0000000..1258ecc
--- /dev/null
+++ b/src/tests/master_maintenance_tests.cpp
@@ -0,0 +1,239 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <string>
+
+#include <process/clock.hpp>
+#include <process/future.hpp>
+#include <process/http.hpp>
+#include <process/pid.hpp>
+#include <process/time.hpp>
+
+#include <stout/duration.hpp>
+#include <stout/json.hpp>
+#include <stout/net.hpp>
+#include <stout/option.hpp>
+#include <stout/strings.hpp>
+#include <stout/stringify.hpp>
+#include <stout/try.hpp>
+
+#include "common/protobuf_utils.hpp"
+
+#include "master/master.hpp"
+
+#include "slave/flags.hpp"
+
+#include "tests/mesos.hpp"
+#include "tests/utils.hpp"
+
+using mesos::internal::master::Master;
+
+using mesos::internal::slave::Slave;
+
+using process::Clock;
+using process::Future;
+using process::PID;
+using process::Time;
+
+using process::http::BadRequest;
+using process::http::OK;
+using process::http::Response;
+
+using mesos::internal::protobuf::maintenance::createSchedule;
+using mesos::internal::protobuf::maintenance::createUnavailability;
+using mesos::internal::protobuf::maintenance::createWindow;
+
+using std::string;
+
+using testing::DoAll;
+
+namespace mesos {
+namespace internal {
+namespace tests {
+
+class MasterMaintenanceTest : public MesosTest
+{
+public:
+  virtual void SetUp()
+  {
+    MesosTest::SetUp();
+
+    // Initialize the default POST header.
+    headers["Content-Type"] = "application/json";
+
+    // Initialize some `MachineID`s.
+    machine1.set_hostname("Machine1");
+    machine2.set_ip("0.0.0.2");
+    machine3.set_hostname("Machine3");
+    machine3.set_ip("0.0.0.3");
+
+    // Initialize the default `Unavailability`.
+    unavailability = createUnavailability(Clock::now());
+  }
+
+
+  // Default headers for all POST's to maintenance endpoints.
+  hashmap<string, string> headers;
+
+  // Some generic `MachineID`s that can be used in this test.
+  MachineID machine1;
+  MachineID machine2;
+  MachineID machine3;
+
+  // Default unavailability.  Used when the test does not care
+  // about the value of the unavailability.
+  Unavailability unavailability;
+};
+
+
+// Posts valid and invalid schedules to the maintenance schedule endpoint.
+TEST_F(MasterMaintenanceTest, UpdateSchedule)
+{
+  // Set up a master.
+  Try<PID<Master>> master = StartMaster();
+  ASSERT_SOME(master);
+
+  // Extra machine used in this test.
+  // It isn't filled in, so it's incorrect.
+  MachineID badMachine;
+
+  // Post a valid schedule with one machine.
+  maintenance::Schedule schedule = createSchedule(
+      {createWindow({machine1}, unavailability)});
+
+  Future<Response> response = process::http::post(
+      master.get(),
+      "maintenance/schedule",
+      headers,
+      stringify(JSON::Protobuf(schedule)));
+
+  AWAIT_EXPECT_RESPONSE_STATUS_EQ(OK().status, response);
+
+  // Get the maintenance schedule.
+  response = process::http::get(
+      master.get(),
+      "maintenance/schedule");
+
+  AWAIT_EXPECT_RESPONSE_STATUS_EQ(OK().status, response);
+
+  // Check that the schedule was saved.
+  Try<JSON::Object> masterSchedule_ =
+    JSON::parse<JSON::Object>(response.get().body);
+
+  ASSERT_SOME(masterSchedule_);
+  Try<mesos::maintenance::Schedule> masterSchedule =
+    ::protobuf::parse<mesos::maintenance::Schedule>(masterSchedule_.get());
+
+  ASSERT_SOME(masterSchedule);
+  ASSERT_EQ(1, masterSchedule.get().windows().size());
+  ASSERT_EQ(1, masterSchedule.get().windows(0).machine_ids().size());
+  ASSERT_EQ(
+      "Machine1",
+      masterSchedule.get().windows(0).machine_ids(0).hostname());
+
+  // Try to replace with an invalid schedule with an empty window.
+  schedule = createSchedule(
+      {createWindow({}, unavailability)});
+
+  response = process::http::post(
+      master.get(),
+      "maintenance/schedule",
+      headers,
+      stringify(JSON::Protobuf(schedule)));
+
+  AWAIT_EXPECT_RESPONSE_STATUS_EQ(BadRequest().status, response);
+
+  // Update the schedule with an unavailability with negative time.
+  schedule = createSchedule({
+      createWindow(
+          {machine1},
+          createUnavailability(Time::create(-10).get()))});
+
+  response = process::http::post(
+      master.get(),
+      "maintenance/schedule",
+      headers,
+      stringify(JSON::Protobuf(schedule)));
+
+  AWAIT_EXPECT_RESPONSE_STATUS_EQ(OK().status, response);
+
+  // Try to replace with an invalid schedule with a negative duration.
+  schedule = createSchedule({
+      createWindow(
+          {machine1},
+          createUnavailability(Clock::now(), Seconds(-10)))});
+
+  response = process::http::post(
+      master.get(),
+      "maintenance/schedule",
+      headers,
+      stringify(JSON::Protobuf(schedule)));
+
+  AWAIT_EXPECT_RESPONSE_STATUS_EQ(BadRequest().status, response);
+
+  // Try to replace with another invalid schedule with duplicate machines.
+  schedule = createSchedule({
+      createWindow({machine1}, unavailability),
+      createWindow({machine1}, unavailability)});
+
+  response = process::http::post(
+      master.get(),
+      "maintenance/schedule",
+      headers,
+      stringify(JSON::Protobuf(schedule)));
+
+  AWAIT_EXPECT_RESPONSE_STATUS_EQ(BadRequest().status, response);
+
+  // Try to replace with an invalid schedule with a badly formed MachineInfo.
+  schedule = createSchedule(
+      {createWindow({badMachine}, unavailability)});
+
+  response = process::http::post(
+      master.get(),
+      "maintenance/schedule",
+      headers,
+      stringify(JSON::Protobuf(schedule)));
+
+  AWAIT_EXPECT_RESPONSE_STATUS_EQ(BadRequest().status, response);
+
+  // Post a valid schedule with two machines.
+  schedule = createSchedule(
+      {createWindow({machine1, machine2}, unavailability)});
+
+  response = process::http::post(
+      master.get(),
+      "maintenance/schedule",
+      headers,
+      stringify(JSON::Protobuf(schedule)));
+
+  AWAIT_EXPECT_RESPONSE_STATUS_EQ(OK().status, response);
+
+  // Delete the schedule (via an empty schedule).
+  schedule = createSchedule({});
+  response = process::http::post(
+      master.get(),
+      "maintenance/schedule",
+      headers,
+      stringify(JSON::Protobuf(schedule)));
+
+  AWAIT_EXPECT_RESPONSE_STATUS_EQ(OK().status, response);
+}
+
+} // namespace tests {
+} // namespace internal {
+} // namespace mesos {

http://git-wip-us.apache.org/repos/asf/mesos/blob/3b0fe5cc/src/tests/registrar_tests.cpp
----------------------------------------------------------------------
diff --git a/src/tests/registrar_tests.cpp b/src/tests/registrar_tests.cpp
index 032e644..567934c 100644
--- a/src/tests/registrar_tests.cpp
+++ b/src/tests/registrar_tests.cpp
@@ -25,6 +25,7 @@
 
 #include <mesos/type_utils.hpp>
 
+#include <process/clock.hpp>
 #include <process/gmock.hpp>
 #include <process/gtest.hpp>
 #include <process/pid.hpp>
@@ -45,6 +46,7 @@
 #include "messages/state.hpp"
 
 #include "master/flags.hpp"
+#include "master/maintenance.hpp"
 #include "master/master.hpp"
 #include "master/registrar.hpp"
 
@@ -68,6 +70,12 @@ using std::set;
 using std::string;
 using std::vector;
 
+using process::Clock;
+
+using mesos::internal::protobuf::maintenance::createSchedule;
+using mesos::internal::protobuf::maintenance::createUnavailability;
+using mesos::internal::protobuf::maintenance::createWindow;
+
 using testing::_;
 using testing::DoAll;
 using testing::Eq;
@@ -81,6 +89,9 @@ namespace mesos {
 namespace internal {
 namespace tests {
 
+using namespace mesos::maintenance;
+using namespace mesos::internal::master::maintenance;
+
 using state::Entry;
 using state::LogStorage;
 using state::Storage;
@@ -316,6 +327,179 @@ TEST_P(RegistrarTest, Remove)
 }
 
 
+// NOTE: For the following tests, the state of the registrar can
+// only be viewed once per instantiation of the registrar.
+// To check the result of each operation, we must re-construct
+// the registrar, which is done by putting the code into scoped blocks.
+
+// TODO(josephw): Consider refactoring these maintenance operation tests
+// to use a helper function for each un-named scoped block.
+// For example:
+//   MaintenanceTest(flags, state, [=](const Registry& registry) {
+//     // Checks and operations.  i.e.:
+//     EXPECT_EQ(1, registry.get().schedules().size());
+//   });
+
+// Adds maintenance schedules to the registry, one machine at a time.
+// Then removes machines from the schedule.
+TEST_P(RegistrarTest, UpdateMaintenanceSchedule)
+{
+  // Machine definitions used in this test.
+  MachineID machine1;
+  machine1.set_ip("0.0.0.1");
+
+  MachineID machine2;
+  machine2.set_hostname("2");
+
+  MachineID machine3;
+  machine3.set_hostname("3");
+  machine3.set_ip("0.0.0.3");
+
+  Unavailability unavailability = createUnavailability(Clock::now());
+
+  {
+    // Prepare the registrar.
+    Registrar registrar(flags, state);
+    AWAIT_READY(registrar.recover(master));
+
+    // Schedule one machine for maintenance.
+    maintenance::Schedule schedule = createSchedule(
+        {createWindow({machine1}, unavailability)});
+
+    AWAIT_READY(registrar.apply(
+        Owned<Operation>(new UpdateSchedule(schedule))));
+  }
+
+  {
+    // Check that one schedule and one machine info was made.
+    Registrar registrar(flags, state);
+    Future<Registry> registry = registrar.recover(master);
+    AWAIT_READY(registry);
+
+    EXPECT_EQ(1, registry.get().schedules().size());
+    EXPECT_EQ(1, registry.get().schedules(0).windows().size());
+    EXPECT_EQ(1, registry.get().schedules(0).windows(0).machine_ids().size());
+    EXPECT_EQ(1, registry.get().machines().machines().size());
+    EXPECT_EQ(
+        MachineInfo::DRAINING,
+        registry.get().machines().machines(0).info().mode());
+
+    // Extend the schedule by one machine (in a different window).
+    maintenance::Schedule schedule = createSchedule({
+        createWindow({machine1}, unavailability),
+        createWindow({machine2}, unavailability)});
+
+    AWAIT_READY(registrar.apply(
+        Owned<Operation>(new UpdateSchedule(schedule))));
+  }
+
+  {
+    // Check that both machines are part of maintenance.
+    Registrar registrar(flags, state);
+    Future<Registry> registry = registrar.recover(master);
+    AWAIT_READY(registry);
+
+    EXPECT_EQ(1, registry.get().schedules().size());
+    EXPECT_EQ(2, registry.get().schedules(0).windows().size());
+    EXPECT_EQ(1, registry.get().schedules(0).windows(0).machine_ids().size());
+    EXPECT_EQ(1, registry.get().schedules(0).windows(1).machine_ids().size());
+    EXPECT_EQ(2, registry.get().machines().machines().size());
+    EXPECT_EQ(
+        MachineInfo::DRAINING,
+        registry.get().machines().machines(1).info().mode());
+
+    // Extend a window by one machine.
+    maintenance::Schedule schedule = createSchedule({
+        createWindow({machine1}, unavailability),
+        createWindow({machine2, machine3}, unavailability)});
+
+    AWAIT_READY(registrar.apply(
+        Owned<Operation>(new UpdateSchedule(schedule))));
+  }
+
+  {
+    // Check that all three machines are included.
+    Registrar registrar(flags, state);
+    Future<Registry> registry = registrar.recover(master);
+    AWAIT_READY(registry);
+
+    EXPECT_EQ(1, registry.get().schedules().size());
+    EXPECT_EQ(2, registry.get().schedules(0).windows().size());
+    EXPECT_EQ(1, registry.get().schedules(0).windows(0).machine_ids().size());
+    EXPECT_EQ(2, registry.get().schedules(0).windows(1).machine_ids().size());
+    EXPECT_EQ(3, registry.get().machines().machines().size());
+    EXPECT_EQ(
+        MachineInfo::DRAINING,
+        registry.get().machines().machines(2).info().mode());
+
+    // Rearrange the schedule into one window.
+    maintenance::Schedule schedule = createSchedule(
+        {createWindow({machine1, machine2, machine3}, unavailability)});
+
+    AWAIT_READY(registrar.apply(
+        Owned<Operation>(new UpdateSchedule(schedule))));
+  }
+
+  {
+    // Check that the machine infos are unchanged, but the schedule is.
+    Registrar registrar(flags, state);
+    Future<Registry> registry = registrar.recover(master);
+    AWAIT_READY(registry);
+
+    EXPECT_EQ(1, registry.get().schedules().size());
+    EXPECT_EQ(1, registry.get().schedules(0).windows().size());
+    EXPECT_EQ(3, registry.get().schedules(0).windows(0).machine_ids().size());
+    EXPECT_EQ(3, registry.get().machines().machines().size());
+    EXPECT_EQ(
+        MachineInfo::DRAINING,
+        registry.get().machines().machines(0).info().mode());
+
+    EXPECT_EQ(
+        MachineInfo::DRAINING,
+        registry.get().machines().machines(1).info().mode());
+
+    EXPECT_EQ(
+        MachineInfo::DRAINING,
+        registry.get().machines().machines(2).info().mode());
+
+    // Delete one machine from the schedule.
+    maintenance::Schedule schedule = createSchedule(
+        {createWindow({machine2, machine3}, unavailability)});
+
+    AWAIT_READY(registrar.apply(
+        Owned<Operation>(new UpdateSchedule(schedule))));
+  }
+
+  {
+    // Check that one machine info is removed.
+    Registrar registrar(flags, state);
+    Future<Registry> registry = registrar.recover(master);
+    AWAIT_READY(registry);
+
+    EXPECT_EQ(1, registry.get().schedules().size());
+    EXPECT_EQ(1, registry.get().schedules(0).windows().size());
+    EXPECT_EQ(2, registry.get().schedules(0).windows(0).machine_ids().size());
+    EXPECT_EQ(2, registry.get().machines().machines().size());
+
+    // Delete all machines from the schedule.
+    maintenance::Schedule schedule;
+    AWAIT_READY(registrar.apply(
+        Owned<Operation>(new UpdateSchedule(schedule))));
+  }
+
+  {
+    // Check that all statuses are removed.
+    Registrar registrar(flags, state);
+    Future<Registry> registry = registrar.recover(master);
+    AWAIT_READY(registry);
+    
+    EXPECT_EQ(1, registry.get().schedules().size());
+    EXPECT_EQ(0, registry.get().schedules(0).windows().size());
+    EXPECT_EQ(0, registry.get().machines().machines().size());
+  }
+}
+
+
 TEST_P(RegistrarTest, Bootstrap)
 {
   // Run 1 readmits a slave that is not present.

Reply via email to