// 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 <cstdint>
#include <functional>
#include <memory>
#include <ostream>
#include <string>
#include <thread>
#include <unordered_map>
#include <utility>
#include <vector>

#include <glog/logging.h>
#include <gtest/gtest.h>

#include "kudu/consensus/consensus.pb.h"
#include "kudu/consensus/metadata.pb.h"
#include "kudu/gutil/map-util.h"
#include "kudu/gutil/port.h"
#include "kudu/gutil/stl_util.h"
#include "kudu/gutil/strings/join.h"
#include "kudu/gutil/strings/substitute.h"
#include "kudu/integration-tests/cluster_itest_util.h"
#include "kudu/integration-tests/cluster_verifier.h"
#include "kudu/integration-tests/external_mini_cluster-itest-base.h"
#include "kudu/integration-tests/test_workload.h"
#include "kudu/master/master.pb.h"
#include "kudu/master/master.proxy.h"
#include "kudu/mini-cluster/external_mini_cluster.h"
#include "kudu/mini-cluster/mini_cluster.h"
#include "kudu/rpc/rpc_controller.h"
#include "kudu/tools/tool_test_util.h"
#include "kudu/util/metrics.h"
#include "kudu/util/monotime.h"
#include "kudu/util/net/sockaddr.h"
#include "kudu/util/path_util.h"
#include "kudu/util/scoped_cleanup.h"
#include "kudu/util/status.h"
#include "kudu/util/test_macros.h"
#include "kudu/util/test_util.h"

METRIC_DECLARE_entity(server);
METRIC_DECLARE_histogram(
    handler_latency_kudu_tserver_TabletCopyService_BeginTabletCopySession);

using kudu::master::ChangeTServerStateRequestPB;
using kudu::master::ChangeTServerStateResponsePB;
using kudu::cluster::ExternalDaemon;
using kudu::cluster::ExternalMiniClusterOptions;
using kudu::cluster::ExternalTabletServer;
using kudu::cluster::LocationInfo;
using kudu::consensus::ConsensusStatePB;
using kudu::consensus::HealthReportPB;
using kudu::consensus::IncludeHealthReport;
using kudu::itest::GetInt64Metric;
using kudu::master::MasterServiceProxy;
using kudu::master::TServerStateChangePB;
using kudu::tools::RunActionPrependStdoutStderr;
using std::function;
using std::pair;
using std::shared_ptr;
using std::string;
using std::thread;
using std::unique_ptr;
using std::unordered_map;
using std::vector;
using strings::Substitute;

namespace kudu {
namespace itest {

typedef pair<unordered_map<string, TServerDetails*>, unique_ptr<ValueDeleter>> MapAndDeleter;
static const vector<string> kTServerFlags = {
  // Set a low unavailability timeout so replicas are considered failed and can
  // be re-replicated more quickly.
  "--raft_heartbeat_interval_ms=100",
  "--follower_unavailable_considered_failed_sec=2",
  // Disable log GC in case our write workloads lead to eviction because
  // consensus will consider replicas that are too far behind unrecoverable
  // and will evict them regardless of maintenance mode.
  "--enable_log_gc=false",
};
static const MonoDelta kDurationForSomeHeartbeats = MonoDelta::FromSeconds(3);

class MaintenanceModeITest : public ExternalMiniClusterITestBase {
 public:
  void SetUpCluster(int num_tservers) {
    ExternalMiniClusterOptions opts;
    opts.num_tablet_servers = num_tservers;
    opts.extra_tserver_flags = kTServerFlags;
    NO_FATALS(StartClusterWithOpts(std::move(opts)));
    const auto& addr = cluster_->master(0)->bound_rpc_addr();
    m_proxy_.reset(new MasterServiceProxy(cluster_->messenger(), addr, addr.host()));
  }

  // Perform the given state change on the given tablet server.
  Status ChangeTServerState(const string& uuid, TServerStateChangePB::StateChange change) {
    ChangeTServerStateRequestPB req;
    ChangeTServerStateResponsePB resp;
    TServerStateChangePB* state_change = req.mutable_change();
    state_change->set_uuid(uuid);
    state_change->set_change(change);
    rpc::RpcController rpc;
    return cluster_->master_proxy()->ChangeTServerState(req, &resp, &rpc);
  }

  // Checks whether tablet copies have started.
  void ExpectStartedTabletCopies(bool should_have_started) {
    bool has_started = false;
    for (int i = 0; i < cluster_->num_tablet_servers(); i++) {
      ExternalTabletServer* tserver = cluster_->tablet_server(i);
      if (tserver->IsShutdown()) {
        continue;
      }
      int64_t copies_started = 0;
      ASSERT_OK(GetInt64Metric(tserver->bound_http_hostport(), &METRIC_ENTITY_server,
          /*entity_id=*/nullptr,
          &METRIC_handler_latency_kudu_tserver_TabletCopyService_BeginTabletCopySession,
          "total_count", &copies_started));
      if (copies_started > 0) {
        has_started = true;
        break;
      }
    }
    ASSERT_EQ(should_have_started, has_started);
  }

  // Return the number of failed replicas there are in the cluster, according
  // to the tablet leaders.
  Status GetReplicaHealthCounts(const unordered_map<string, TServerDetails*>& ts_map,
                                int* num_replicas_failed, int* num_replicas_healthy) {
    int num_failed = 0;
    int num_healthy = 0;
    for (int i = 0; i < cluster_->num_tablet_servers(); i++) {
      ExternalTabletServer* tserver = cluster_->tablet_server(i);
      if (tserver->IsShutdown()) {
        continue;
      }
      const string& uuid = tserver->uuid();
      const TServerDetails* ts_details = FindOrDie(ts_map, uuid);
      vector<string> tablet_ids;
      RETURN_NOT_OK(ListRunningTabletIds(ts_details, MonoDelta::FromSeconds(30), &tablet_ids));
      for (const auto& tablet_id : tablet_ids) {
        ConsensusStatePB consensus_state;
        RETURN_NOT_OK(GetConsensusState(ts_details, tablet_id, MonoDelta::FromSeconds(30),
            IncludeHealthReport::INCLUDE_HEALTH_REPORT, &consensus_state));
        // Only consider the health states reported by the leaders.
        if (consensus_state.leader_uuid() != uuid) {
          continue;
        }
        // Go through all the peers and tally up any that are failed.
        const auto& committed_config = consensus_state.committed_config();
        for (int p = 0; p < committed_config.peers_size(); p++) {
          const auto& peer = committed_config.peers(p);
          if (peer.has_health_report()) {
            switch (peer.health_report().overall_health()) {
              case HealthReportPB::FAILED:
                num_failed++;
                break;
              case HealthReportPB::HEALTHY:
                num_healthy++;
                break;
              default:
                continue;
            }
          }
        }
      }
    }
    *num_replicas_failed = num_failed;
    *num_replicas_healthy = num_healthy;
    return Status::OK();
  }

  // Asserts that the ts_map will eventually have the given number of failed and healthy replicas
  // across the cluster.
  void AssertEventuallyHealthCounts(const unordered_map<string, TServerDetails*>& ts_map,
                                    int expected_failed, int expected_healthy) {
    ASSERT_EVENTUALLY([&] {
      int num_failed;
      int num_healthy;
      ASSERT_OK(GetReplicaHealthCounts(ts_map, &num_failed, &num_healthy));
      ASSERT_EQ(expected_failed, num_failed);
      ASSERT_EQ(expected_healthy, num_healthy);
    });
  }

  void GenerateTServerMap(MapAndDeleter* map_and_deleter) {
    unordered_map<string, TServerDetails*> ts_map;
    ASSERT_EVENTUALLY([&] {
      ASSERT_OK(CreateTabletServerMap(m_proxy_, cluster_->messenger(), &ts_map));
      auto cleanup = MakeScopedCleanup([&] {
        STLDeleteValues(&ts_map);
      });
      ASSERT_EQ(cluster_->num_tablet_servers(), ts_map.size());
      cleanup.cancel();
    });
    map_and_deleter->first = std::move(ts_map);
    map_and_deleter->second.reset(new ValueDeleter(&map_and_deleter->first));
  }

 protected:
  shared_ptr<MasterServiceProxy> m_proxy_;
};

class MaintenanceModeRF3ITest : public MaintenanceModeITest {
 public:
  void SetUp() override {
    SKIP_IF_SLOW_NOT_ALLOWED();
    NO_FATALS(MaintenanceModeITest::SetUp());
    NO_FATALS(SetUpCluster(3));
  }
  void TearDown() override {
    SKIP_IF_SLOW_NOT_ALLOWED();
    NO_FATALS(MaintenanceModeITest::TearDown());
  }
};

// Test that placing a tablet server in maintenance mode leads to the failed
// replicas on that server not being re-replicated.
TEST_F(MaintenanceModeRF3ITest, TestFailedTServerInMaintenanceModeDoesntRereplicate) {
  // This test will sleep a bit to ensure the master has had some time to
  // receive heartbeats.
  SKIP_IF_SLOW_NOT_ALLOWED();
  const int kNumTablets = 6;
  const int total_replicas = kNumTablets * 3;

  // Create the table with three tablet servers and then add one so we're
  // guaranteed that the replicas are all on the first three servers.
  // Restarting our only master may lead to network errors and timeouts, but
  // that shouldn't matter w.r.t maintenance mode.
  TestWorkload create_table(cluster_.get());
  create_table.set_num_tablets(kNumTablets);
  create_table.set_network_error_allowed(true);
  create_table.set_timeout_allowed(true);
  create_table.Setup();
  create_table.Start();
  // Add a server so there's one we could move to after bringing down a
  // tserver.
  ASSERT_OK(cluster_->AddTabletServer());
  MapAndDeleter ts_map_and_deleter;
  NO_FATALS(GenerateTServerMap(&ts_map_and_deleter));
  const auto& ts_map = ts_map_and_deleter.first;

  // Do a sanity check that all our replicas are healthy.
  NO_FATALS(AssertEventuallyHealthCounts(ts_map, 0, total_replicas));

  // Put one of the servers in maintenance mode.
  ExternalTabletServer* maintenance_ts = cluster_->tablet_server(0);
  const string maintenance_uuid = maintenance_ts->uuid();
  ASSERT_OK(ChangeTServerState(maintenance_uuid, TServerStateChangePB::ENTER_MAINTENANCE_MODE));

  // Bringing the tablet server down shouldn't lead to re-replication.
  //
  // Note: it's possible to set up re-replication scenarios in other ways (e.g.
  // by hitting an IO error, or falling too far behind when replicating); these
  // should all be treated the same way by virtue of them all using the same
  // health reporting.
  NO_FATALS(maintenance_ts->Shutdown());

  // Now wait a bit for this failure to make its way to the master. The failure
  // shouldn't have led to any re-replication.
  SleepFor(kDurationForSomeHeartbeats);
  NO_FATALS(AssertEventuallyHealthCounts(ts_map, kNumTablets, total_replicas - kNumTablets));
  NO_FATALS(ExpectStartedTabletCopies(/*should_have_started*/false));

  // Restarting the masters shouldn't lead to re-replication either, even
  // though the tablet server is still down.
  NO_FATALS(cluster_->ShutdownNodes(cluster::ClusterNodes::MASTERS_ONLY));
  ASSERT_OK(cluster_->master()->Restart());
  SleepFor(kDurationForSomeHeartbeats);
  NO_FATALS(ExpectStartedTabletCopies(/*should_have_started*/false));

  // Now bring the server back up and wait for it to become healthy. It should
  // be able to do this without tablet copies.
  ASSERT_OK(maintenance_ts->Restart());
  NO_FATALS(AssertEventuallyHealthCounts(ts_map, 0, total_replicas));

  // Since our server is healthy, leaving maintenance mode shouldn't trigger
  // any re-replication either.
  ASSERT_OK(ChangeTServerState(maintenance_uuid, TServerStateChangePB::EXIT_MAINTENANCE_MODE));
  SleepFor(kDurationForSomeHeartbeats);
  NO_FATALS(ExpectStartedTabletCopies(/*should_have_started*/false));

  // Now set maintenance mode, bring the tablet server down, and then exit
  // maintenance mode without bringing the tablet server back up. This should
  // result in tablet copies.
  ASSERT_OK(ChangeTServerState(maintenance_uuid, TServerStateChangePB::ENTER_MAINTENANCE_MODE));
  NO_FATALS(maintenance_ts->Shutdown());
  NO_FATALS(AssertEventuallyHealthCounts(ts_map, kNumTablets, total_replicas - kNumTablets));
  ASSERT_OK(ChangeTServerState(maintenance_uuid, TServerStateChangePB::EXIT_MAINTENANCE_MODE));
  ASSERT_EVENTUALLY([&] {
    NO_FATALS(ExpectStartedTabletCopies(/*should_have_started*/true));
  });

  // All the while, our workload should not have been interrupted. Assert
  // eventually to wait for the rows to converge.
  NO_FATALS(create_table.StopAndJoin());
  ClusterVerifier v(cluster_.get());
  ASSERT_EVENTUALLY([&] {
    NO_FATALS(v.CheckRowCount(create_table.table_name(),
                              ClusterVerifier::EXACTLY, create_table.rows_inserted()));
  });
}

TEST_F(MaintenanceModeRF3ITest, TestMaintenanceModeDoesntObstructMove) {
  SKIP_IF_SLOW_NOT_ALLOWED();
  TestWorkload create_table(cluster_.get());
  create_table.set_num_tablets(1);
  create_table.Setup();

  // Add a tablet server.
  ASSERT_OK(cluster_->AddTabletServer());
  const string& added_uuid = cluster_->tablet_server(3)->uuid();
  MapAndDeleter ts_map_and_deleter;
  NO_FATALS(GenerateTServerMap(&ts_map_and_deleter));
  const auto& ts_map = ts_map_and_deleter.first;

  // Put a tablet server into maintenance mode.
  const string maintenance_uuid = cluster_->tablet_server(0)->uuid();
  const TServerDetails* maintenance_details = FindOrDie(ts_map, maintenance_uuid);
  vector<string> mnt_tablet_ids;
  ASSERT_OK(ListRunningTabletIds(maintenance_details, MonoDelta::FromSeconds(30), &mnt_tablet_ids));
  ASSERT_EQ(1, mnt_tablet_ids.size());

  ASSERT_OK(ChangeTServerState(maintenance_uuid, TServerStateChangePB::ENTER_MAINTENANCE_MODE));

  // While the maintenance mode tserver is still online, move a tablet from it.
  // This should succeed, because maintenance mode will not obstruct manual
  // movement of replicas.
  ASSERT_OK(RunActionPrependStdoutStderr(Substitute(
      "tablet change_config move_replica $0 $1 $2 $3",
      cluster_->master()->bound_rpc_addr().ToString(),
      mnt_tablet_ids[0],
      maintenance_uuid,
      added_uuid)));
  const TServerDetails* added_details = FindOrDie(ts_map, added_uuid);
  ASSERT_EVENTUALLY([&] {
    vector<string> added_tablet_ids;
    ASSERT_OK(ListRunningTabletIds(added_details, MonoDelta::FromSeconds(30), &added_tablet_ids));
    ASSERT_EQ(1, added_tablet_ids.size());
  });
}

// Test that the health state FAILED_UNRECOVERABLE (e.g. if there's a disk
// error, or if a replica is lagging too much) is still re-replicated during
// maintenance mode.
TEST_F(MaintenanceModeRF3ITest, TestMaintenanceModeDoesntObstructFailedUnrecoverable) {
  SKIP_IF_SLOW_NOT_ALLOWED();
  TestWorkload create_table(cluster_.get());
  create_table.set_num_tablets(1);
  create_table.Setup();
  create_table.Start();

  // Add a tablet server.
  ASSERT_OK(cluster_->AddTabletServer());
  const string& added_uuid = cluster_->tablet_server(3)->uuid();
  MapAndDeleter ts_map_and_deleter;
  NO_FATALS(GenerateTServerMap(&ts_map_and_deleter));
  const auto& ts_map = ts_map_and_deleter.first;

  // Put a tablet server into maintenance mode.
  ExternalDaemon* maintenance_ts = cluster_->tablet_server(0);
  const string maintenance_uuid = maintenance_ts->uuid();
  const TServerDetails* maintenance_details = FindOrDie(ts_map, maintenance_uuid);
  vector<string> mnt_tablet_ids;
  ASSERT_OK(ListRunningTabletIds(maintenance_details, MonoDelta::FromSeconds(30), &mnt_tablet_ids));
  ASSERT_OK(ChangeTServerState(maintenance_uuid, TServerStateChangePB::ENTER_MAINTENANCE_MODE));

  // Now fail the tablet on the server in maintenance mode by injecting a disk
  // error. Also speed up flushes so we actually hit an IO error.
  ASSERT_OK(cluster_->SetFlag(maintenance_ts, "flush_threshold_secs", "1"));
  ASSERT_OK(cluster_->SetFlag(maintenance_ts, "env_inject_eio_globs",
      JoinPathSegments(maintenance_ts->data_dirs()[0], "**")));
  ASSERT_OK(cluster_->SetFlag(maintenance_ts, "env_inject_eio", "1"));

  // Eventually the disk failure will be noted and a copy will be made at the
  // added server.
  const TServerDetails* added_details = FindOrDie(ts_map, added_uuid);
  ASSERT_EVENTUALLY([&] {
    vector<string> added_tablet_ids;
    ASSERT_OK(ListRunningTabletIds(added_details, MonoDelta::FromSeconds(30), &added_tablet_ids));
    ASSERT_EQ(1, added_tablet_ids.size());
  });
  NO_FATALS(create_table.StopAndJoin());
  ClusterVerifier v(cluster_.get());
  ASSERT_EVENTUALLY([&] {
    NO_FATALS(v.CheckRowCount(create_table.table_name(),
                              ClusterVerifier::EXACTLY, create_table.rows_inserted()));
  });
}

class MaintenanceModeRF5ITest : public MaintenanceModeITest {
 public:
  void SetUp() override {
    SKIP_IF_SLOW_NOT_ALLOWED();
    NO_FATALS(MaintenanceModeITest::SetUp());
    NO_FATALS(SetUpCluster(5));
  }
  void TearDown() override {
    SKIP_IF_SLOW_NOT_ALLOWED();
    NO_FATALS(MaintenanceModeITest::TearDown());
  }
};

// Test that a table with RF=5 will still be available through the failure of
// two nodes if one is put in maintenance mode.
TEST_F(MaintenanceModeRF5ITest, TestBackgroundFailureDuringMaintenanceMode) {
  SKIP_IF_SLOW_NOT_ALLOWED();
  // Create some tables with RF=5.
  const int kNumTablets = 3;
  const int total_replicas = kNumTablets * 5;
  TestWorkload create_table(cluster_.get());
  create_table.set_num_tablets(kNumTablets);
  create_table.set_num_replicas(5);
  create_table.Setup();
  create_table.Start();

  // Add a server so we have one empty server to replicate to.
  ASSERT_OK(cluster_->AddTabletServer());
  MapAndDeleter ts_map_and_deleter;
  NO_FATALS(GenerateTServerMap(&ts_map_and_deleter));
  const auto& ts_map = ts_map_and_deleter.first;

  // Do a sanity check that all our replicas are healthy.
  NO_FATALS(AssertEventuallyHealthCounts(ts_map, 0, total_replicas));

  // Enter maintenance mode on a tserver and shut it down.
  ExternalTabletServer* maintenance_ts = cluster_->tablet_server(0);
  const string maintenance_uuid = maintenance_ts->uuid();
  ASSERT_OK(ChangeTServerState(maintenance_uuid, TServerStateChangePB::ENTER_MAINTENANCE_MODE));
  NO_FATALS(maintenance_ts->Shutdown());
  SleepFor(kDurationForSomeHeartbeats);

  // Wait for the failure to be recognized by the other replicas.
  NO_FATALS(AssertEventuallyHealthCounts(ts_map, kNumTablets, total_replicas - kNumTablets));
  NO_FATALS(ExpectStartedTabletCopies(/*should_have_started*/false));

  // Now kill another server. We should be able to see some copies.
  NO_FATALS(cluster_->tablet_server(1)->Shutdown());
  ASSERT_EVENTUALLY([&] {
    NO_FATALS(ExpectStartedTabletCopies(/*should_have_started*/true));
  });
  // Eventually we'll be left with just the unresponsive maintenance mode
  // failed replicas.
  NO_FATALS(AssertEventuallyHealthCounts(ts_map, kNumTablets, total_replicas - kNumTablets));
  // The previously empty tablet server should hold all the replicas that were
  // re-replicated.
  const TServerDetails* added_details = FindOrDie(ts_map, cluster_->tablet_server(5)->uuid());
  ASSERT_EVENTUALLY([&] {
    vector<string> added_tablet_ids;
    ASSERT_OK(ListRunningTabletIds(added_details, MonoDelta::FromSeconds(30), &added_tablet_ids));
    ASSERT_EQ(kNumTablets, added_tablet_ids.size());
  });
  // Now exit maintenance mode and restart the maintenance tserver. The
  // original tablets on the maintenance mode tserver should still exist.
  ASSERT_OK(ChangeTServerState(maintenance_uuid, TServerStateChangePB::EXIT_MAINTENANCE_MODE));
  ASSERT_OK(maintenance_ts->Restart());
  const TServerDetails* maintenance_details = FindOrDie(ts_map, maintenance_uuid);
  vector<string> mnt_tablet_ids;
  ASSERT_EVENTUALLY([&] {
    ASSERT_OK(ListRunningTabletIds(
        maintenance_details, MonoDelta::FromSeconds(30), &mnt_tablet_ids));
    ASSERT_EQ(kNumTablets, mnt_tablet_ids.size());
  });

  // All the while, our workload should not have been interrupted. Assert
  // eventually to wait for the rows to converge.
  NO_FATALS(create_table.StopAndJoin());
  ClusterVerifier v(cluster_.get());
  ASSERT_EVENTUALLY([&] {
    NO_FATALS(v.CheckRowCount(create_table.table_name(),
                              ClusterVerifier::EXACTLY, create_table.rows_inserted()));
  });
}

namespace {

// Performs the given tasks in parallel, returning an error if any of them
// return a non-OK Status.
Status DoInParallel(vector<function<Status()>> tasks,
                    const string& task_description) {
  int num_in_parallel = tasks.size();
  vector<Status> results(num_in_parallel);
  vector<thread> threads;
  for (int i = 0; i < num_in_parallel; i++) {
    threads.emplace_back([&, i] {
      Status s = tasks[i]();
      if (PREDICT_FALSE(!s.ok())) {
        results[i] = s;
      }
    });
  }
  for (auto& t : threads) {
    t.join();
  }
  bool has_errors = false;
  for (int i = 0; i < num_in_parallel; i++) {
    const auto& s = results[i];
    if (!s.ok()) {
      LOG(ERROR) << s.ToString();
      has_errors = true;
    }
  }
  if (has_errors) {
    return Status::IllegalState(
        Substitute("errors while running $0 $1 tasks in parallel",
                    num_in_parallel, task_description));
  }
  return Status::OK();
}

// Repeats 'cmd' until it succeeds, with the given timeout and interval.
Status RunUntilSuccess(const string& cmd, int timeout_secs, int repeat_interval_secs) {
  Status s;
  const MonoTime deadline = MonoTime::Now() + MonoDelta::FromSeconds(timeout_secs);
  const MonoDelta retry_interval = MonoDelta::FromSeconds(repeat_interval_secs);
  while (true) {
    s = RunActionPrependStdoutStderr(cmd);
    if (s.ok()) {
      return Status::OK();
    }
    if (MonoTime::Now() + retry_interval > deadline) {
      return Status::TimedOut(
          Substitute("Running '$0' did not succeed in $1 seconds: $2",
                      cmd, timeout_secs, s.ToString()));
    }
    SleepFor(retry_interval);
  }
  return Status::OK();
}

// Generates locations such that there are a total of 'num_tservers' tablet
// servers spread as evenly as possible across 'num_locs' locations.
ExternalMiniClusterOptions GenerateOpts(int num_tservers, int num_locs) {
  int tservers_per_loc = num_tservers / num_locs;
  int tservers_with_one_more = num_tservers % num_locs;
  CHECK_LT(0, tservers_per_loc);
  LocationInfo locations;
  for (int l = 0; l < num_locs; l++) {
    string loc = Substitute("/L$0", l);
    EmplaceOrDie(&locations, Substitute("/L$0", l),
        tservers_with_one_more < l ? tservers_per_loc + 1 : tservers_per_loc);
  }

  ExternalMiniClusterOptions opts;
  opts.num_tablet_servers = num_tservers;
  opts.extra_master_flags = {
    // Don't bother assigning locations to clients; it's hard to do that
    // correctly with external mini clusters.
    "--master_client_location_assignment_enabled=false",
  };
  opts.extra_tserver_flags = {
    // To speed up leadership transfers from quiescing, let's make heartbeats
    // more frequent.
    "--raft_heartbeat_interval_ms=100",
    // Let's emulate long-running scans by setting a low scan batch size.
    "--scanner_default_batch_size_bytes=100",
  };
  opts.location_info = std::move(locations);
  return opts;
}

struct RollingRestartTestArgs {
  // Cluster opts.
  ExternalMiniClusterOptions opts;
  // Upper bound on the number of tablet servers to restart at the same time.
  int batch_size;
  // Replication to use for the test workload.
  int num_replicas;
  // Whether the rolling restart should fail.
  bool restart_fails;
};

// Convenience builder for more readable test composition.
class ArgsBuilder {
 public:
  ArgsBuilder& batch_size(int batch_size) {
    args_.batch_size = batch_size;
    return *this;
  }
  ArgsBuilder& num_locations(int num_locations) {
    num_locations_ = num_locations;
    return *this;
  }
  ArgsBuilder& num_replicas(int num_replicas) {
    args_.num_replicas = num_replicas;
    return *this;
  }
  ArgsBuilder& num_tservers(int num_tservers) {
    num_tservers_ = num_tservers;
    return *this;
  }
  ArgsBuilder& restart_fails(bool fails) {
    args_.restart_fails = fails;
    return *this;
  }
  RollingRestartTestArgs build() {
    args_.opts = GenerateOpts(num_tservers_, num_locations_);
    return args_;
  }
 private:
  RollingRestartTestArgs args_;
  int num_tservers_ = 0;
  int num_locations_ = 0;
};

} // anonymous namespace

class RollingRestartITest : public MaintenanceModeITest,
                            public ::testing::WithParamInterface<RollingRestartTestArgs> {
 public:
  void SetUp() override {
    SKIP_IF_SLOW_NOT_ALLOWED();
    NO_FATALS(MaintenanceModeITest::SetUp());
    const auto& args = GetParam();
    ExternalMiniClusterOptions opts = args.opts;
    NO_FATALS(StartClusterWithOpts(std::move(opts)));
    const auto& addr = cluster_->master(0)->bound_rpc_addr();
    m_proxy_.reset(new MasterServiceProxy(cluster_->messenger(), addr, addr.host()));
    NO_FATALS(GenerateTServerMap(&ts_map_and_deleter_));
  }
  void TearDown() override {
    SKIP_IF_SLOW_NOT_ALLOWED();
    NO_FATALS(MaintenanceModeITest::TearDown());
  }

  // Create a read-write workload that doesn't use a fault-tolerant scanner.
  unique_ptr<TestWorkload> CreateFaultIntolerantRWWorkload(int num_replicas) {
    unique_ptr<TestWorkload> rw_workload(new TestWorkload(cluster_.get()));
    rw_workload->set_scanner_fault_tolerant(false);
    rw_workload->set_num_replicas(num_replicas);
    rw_workload->set_num_read_threads(3);
    rw_workload->set_num_write_threads(3);
    rw_workload->set_verify_num_rows(false);
    return rw_workload;
  }

  // Returns a list of batches of tablet server UUIDs to restart in parallel,
  // each of size at most 'batch_size'. Batches are generated within each
  // location.
  vector<vector<string>> GetRestartBatches(int batch_size) {
    unordered_map<string, vector<string>> cur_tservers_by_loc;
    vector<vector<string>> restart_batches;
    for (const auto& ts_and_details : ts_map_and_deleter_.first) {
      const auto& uuid = ts_and_details.first;
      const auto& loc = ts_and_details.second->location;
      auto& cur_batch = LookupOrInsert(&cur_tservers_by_loc, loc, {});
      cur_batch.emplace_back(uuid);
      // If we've reached our desired batch size for this location, put it in
      // the return set.
      if (cur_batch.size() >= batch_size) {
        restart_batches.emplace_back(std::move(cur_batch));
        CHECK_EQ(1, cur_tservers_by_loc.erase(loc));
      }
    }
    // Create batches out of the remaining, suboptimally-sized batches.
    for (auto& ts_and_batch : cur_tservers_by_loc) {
      restart_batches.emplace_back(std::move(ts_and_batch.second));
    }
    return restart_batches;
  }

  // Takes the list of tablet servers and restarts them in parallel in such a
  // way that shouldn't affect on-going workloads.
  Status RollingRestartTServers(const vector<string>& ts_to_restart) {
    // Begin maintenance mode on the servers so we'll stop assigning replicas
    // to them, and begin quiescing them so the server itself stops accepting
    // more work.
    vector<function<Status()>> setup_rr_tasks;
    for (const auto& ts_id : ts_to_restart) {
      setup_rr_tasks.emplace_back([&, ts_id] {
        RETURN_NOT_OK(RunActionPrependStdoutStderr(
            Substitute("tserver state enter_maintenance $0 $1",
                       cluster_->master()->bound_rpc_addr().ToString(), ts_id)));
        return RunUntilSuccess(
            Substitute("tserver quiesce start $0 --error_if_not_fully_quiesced",
                       cluster_->tablet_server_by_uuid(ts_id)->bound_rpc_addr().ToString()),
            /*timeout_secs*/30, /*repeat_interval_secs*/1);
      });
    }
    RETURN_NOT_OK(DoInParallel(std::move(setup_rr_tasks), "rolling restart setup"));

    // Restart the tservers. Note: we don't do this in parallel because
    // wrangling multiple processes from different threads is messy.
    for (const auto& ts_id : ts_to_restart) {
      ExternalTabletServer* ts = cluster_->tablet_server_by_uuid(ts_id);
      ts->Shutdown();
      RETURN_NOT_OK(ts->Restart());
    }

    // Wait for ksck to become healthy.
    RETURN_NOT_OK_PREPEND(RunUntilSuccess(
        Substitute("cluster ksck $0", cluster_->master()->bound_rpc_addr().ToString()),
        /*timeout_secs*/60, /*repeat_interval_secs*/5),
        "cluster didn't become healthy");

    // Now clean up persistent state.
    vector<function<Status()>> cleanup_tasks;
    for (const auto& ts_id : ts_to_restart) {
      cleanup_tasks.emplace_back([&, ts_id] {
        return RunActionPrependStdoutStderr(
            Substitute("tserver state exit_maintenance $0 $1",
                       cluster_->master()->bound_rpc_addr().ToString(),
                       ts_id));

      });
    }
    return DoInParallel(std::move(cleanup_tasks), "rolling restart cleanup");
  }

 protected:
  MapAndDeleter ts_map_and_deleter_;
};

TEST_P(RollingRestartITest, TestWorkloads) {
  SKIP_IF_SLOW_NOT_ALLOWED();
  const auto& args = GetParam();
  unique_ptr<TestWorkload> rw = CreateFaultIntolerantRWWorkload(args.num_replicas);
  rw->set_read_timeout_millis(10000);
  // If we're expecting the rolling restart to fail, e.g. because we can't
  // fully quiesce our servers, chances are our workload can't complete either
  // because too many servers are quiescing and none can serve scans or writes.
  rw->set_read_errors_allowed(args.restart_fails);
  rw->set_timeout_allowed(args.restart_fails);
  rw->Setup();
  rw->Start();
  vector<vector<string>> restart_batches = GetRestartBatches(args.batch_size);
  for (const auto& batch : restart_batches) {
    LOG(INFO) << Substitute("Restarting batch of $0 tservers: $1",
                            batch.size(), JoinStrings(batch, ","));
    if (args.restart_fails) {
      ASSERT_EVENTUALLY([&] {
        Status s = RollingRestartTServers(batch);
        ASSERT_FALSE(s.ok());
      });
    } else {
      ASSERT_OK(RollingRestartTServers(batch));
    }
  }
  NO_FATALS(rw->StopAndJoin());
  if (args.restart_fails) {
    ASSERT_FALSE(rw->read_errors().empty());
  }
}

INSTANTIATE_TEST_SUITE_P(RollingRestartArgs, RollingRestartITest, ::testing::Values(
    // Basic RF=3 case.
    ArgsBuilder().num_tservers(4)
                 .num_locations(1)
                 .batch_size(1)
                 .num_replicas(3)
                 .restart_fails(false)
                 .build(),
    // Basic RF=5 case. The larger replication factor lets us increase the
    // restart batch size.
    ArgsBuilder().num_tservers(6)
                 .num_locations(1)
                 .batch_size(2)
                 .num_replicas(5)
                 .restart_fails(false)
                 .build(),
    // RF=3 case with location awareness. The location awareness lets us
    // increase the restart batch size.
    ArgsBuilder().num_tservers(6)
                 .num_locations(3)
                 .batch_size(2)
                 .num_replicas(3)
                 .restart_fails(false)
                 .build(),
    // RF=3 case with location awareness, but with an even larger batch size.
    ArgsBuilder().num_tservers(9)
                 .num_locations(3)
                 .batch_size(3)
                 .num_replicas(3)
                 .restart_fails(false)
                 .build(),
    // Basic RF=3 case, but with too large a batch size. With too large a batch
    // size, the tablet servers won't be able to fully relinquish leadership
    // and quiesce; the restart process should thus fail.
    ArgsBuilder().num_tservers(4)
                 .num_locations(1)
                 .batch_size(4)
                 .num_replicas(3)
                 .restart_fails(true)
                 .build(),
    // The same goes for the case with a single replica.
    ArgsBuilder().num_tservers(1)
                 .num_locations(1)
                 .batch_size(1)
                 .num_replicas(1)
                 .restart_fails(true)
                 .build()
));

} // namespace itest
} // namespace kudu