<gh_stars>0
/**
 * Copyright (c) 2017-present, Facebook, Inc. and its affiliates.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */
#include "logdevice/common/WeightAwareNodeSetSelector.h"

#include <map>
#include <queue>
#include <utility>

#include <folly/Random.h>

#include "logdevice/common/configuration/nodes/utils.h"
#include "logdevice/common/hash.h"

namespace facebook { namespace logdevice {

NodeSetSelector::Result WeightAwareNodeSetSelector::getStorageSet(
    logid_t log_id,
    const Configuration* cfg,
    const configuration::nodes::NodesConfiguration& nodes_configuration,
    nodeset_size_t target_nodeset_size,
    uint64_t seed,
    const EpochMetaData* prev,
    const Options* options) {
  Result res;
  res.decision = Decision::FAILED;

  auto logcfg = cfg->getLogGroupByIDShared(log_id);
  if (!logcfg) {
    return res;
  }

  ReplicationProperty replication_property =
      ReplicationProperty::fromLogAttributes(logcfg->attrs());

  res.signature = hash_tuple({15345803578954886993ul, // random salt
                              consistentHashing_,
                              seed,
                              nodes_configuration.getStorageNodesHash(),
                              target_nodeset_size});
  if (prev != nullptr && prev->nodeset_params.signature == res.signature &&
      prev->replication == replication_property) {
    res.decision = Decision::KEEP;
    return res;
  }

  // Nodeset will span all domains of this scope.
  NodeLocationScope replication_scope;
  // We'll try to pick at least this many nodes from each domain of scope
  // replication_scope.
  int min_nodes_per_domain;

  auto replication_factors =
      replication_property.getDistinctReplicationFactors();
  if (replication_factors[0].first <= NodeLocationScope::NODE) {
    // No cross-domain replication and no sequencer locality.
    replication_scope = NodeLocationScope::ROOT;
    min_nodes_per_domain = replication_factors[0].second;
  } else {
    replication_scope = replication_factors[0].first;
    // If some domain is smaller than this, the current implementation of
    // WeightedCopySetSelector will underutilize that domain.
    min_nodes_per_domain =
        replication_factors.back().second - replication_factors[0].second + 1;
  }

  struct CandidateNode {
    uint64_t shard_id_hash;
    ShardID shard_id;
    bool is_writable;
  };

  struct Domain {
    int num_picked = 0;
    int num_picked_writable = 0;

    uint64_t priority;

    // The list of nodes sorted by hash.
    // To pick a node, we pop one from the back.
    std::vector<CandidateNode> nodes;
  };
  std::map<std::string, Domain> domains;

  const auto& membership = nodes_configuration.getStorageMembership();
  for (const auto node : *membership) {
    // Filter nodes excluded from `options`.
    if (options != nullptr && options->exclude_nodes.count(node)) {
      continue;
    }

    const auto num_shards = nodes_configuration.getNumShards(node);
    ld_check(num_shards > 0);
    shard_index_t shard_idx = mapLogToShard_(log_id, num_shards);
    ShardID shard = ShardID(node, shard_idx);

    // Filter nodes that shouldn't be included in nodesets
    if (!configuration::nodes::shouldIncludeInNodesetSelection(
            nodes_configuration, shard)) {
      continue;
    }

    const auto* sd = nodes_configuration.getNodeServiceDiscovery(node);
    ld_check(sd != nullptr);

    std::string location_str;
    if (replication_scope == NodeLocationScope::ROOT) {
      // All nodes are in the same replication domain.
    } else {
      if (!sd->location.hasValue()) {
        ld_error("Can't select nodeset because node %d (%s) does not have "
                 "location information",
                 node,
                 sd->address.toString().c_str());
        return res;
      }

      const NodeLocation& location = sd->location.value();
      assert(!location.isEmpty());
      if (!location.scopeSpecified(replication_scope)) {
        ld_error("Can't select nodeset because location %s of node %d (%s) "
                 "doesn't have location for scope %s.",
                 location.toString().c_str(),
                 node,
                 sd->address.toString().c_str(),
                 NodeLocation::scopeNames()[replication_scope].c_str());
        return res;
      }
      location_str = location.getDomain(replication_scope, node);
    }

    CandidateNode n;
    n.shard_id = shard;
    n.is_writable = membership->canWriteToShard(shard);

    // If consistent hashing is toggled, hash the log id along with the shard
    // ID, thereby allowing us to create a unique, deterministic ranking of
    // Shard ID's for each log. Else, hash the log id along with a random
    // number, which is equivalent to randomly sorting the Shard ID's for
    // each log.
    if (consistentHashing_) {
      n.shard_id_hash =
          hash_tuple({seed, log_id.val(), ShardID::Hash()(n.shard_id)});
    } else {
      n.shard_id_hash = folly::Random::rand64();
    }
    domains[location_str].nodes.push_back(n);
  }

  for (auto& kv : domains) {
    Domain* d = &kv.second;
    d->priority = consistentHashing_
        ? hash_tuple({seed, log_id.val(), folly::hash::fnv64(kv.first)})
        : folly::Random::rand64();
    std::sort(d->nodes.begin(),
              d->nodes.end(),
              [](const CandidateNode& a, const CandidateNode& b) {
                // Sort in order of decreasing hash, so that we pick nodes with
                // smaller hash first (for historical reasons).
                return std::tie(a.shard_id_hash, a.shard_id) >
                    std::tie(b.shard_id_hash, b.shard_id);
              });
  }

  // Form the nodeset by repeatedly popping entries from the shuffled node list
  // of the domain with the smallest number of picked nodes. If there's a tie,
  // take a random domain. This way the nodeset is distributed across domains as
  // evenly as possible. Keep picking nodes until the nodeset is "good", i.e.
  // big enough and has enough nodes in each available domain.
  //
  // There's a subtlety about whether or not to allow picking temporarily
  // disabled ("unwritable") nodes. On the one hand, we want to pick disabled
  // nodes because, if a full domain gets temporarily disabled and then comes
  // back, we want the domain to be included in historical nodesets, so that it
  // can receive writes from rebuilding and recovery. On the other hand, if lots
  // of nodes are disabled, we don't want to end up with a nodeset that consists
  // mostly of unwritable nodes and doesn't have enough writable nodes to
  // replicate new records.
  //
  // So we do a bit of both: we build a nodeset that is "good" is we count
  // only writable nodes and also "good" if we count both writable and
  // unwritable nodes. Moreover, the writable part of the nodeset has the same
  // distribution as if the unwritable nodes didn't exist; this ensures that
  // the data is evenly distributed among the writable nodes. The resulting
  // nodeset is at most twice as big as it would be without this whole trick.
  //
  // This is done by building the nodeset in two stages:
  //  1. Allow picking both writable and unwritable nodes. Pick nodes until
  //     the nodeset is "good".
  //  2. Continue picking nodes but allow picking only writable nodes, and
  //     change the comparator to pick domain with the smallest number of
  //     _writable_ nodes picked so far. Do it until the writable part of the
  //     nodeset is good.
  //
  // Note that, although the data distribution is correct for the current state
  // of the nodes, it may become biased when the unwritable nodes become
  // writable again. Consider this example. Replication factor is 1, nodeset
  // size is 1, there are 2 nodes: N0 is writable, N1 is unwritable.
  // We'll select nodeset {N0} for ~half the logs, {N0, N1} for the other ~half.
  // When N1 becomes writable, it'll receive only 25% of writes.

  auto select_nodes = [&](bool only_writable) {
    auto cmp = [&](Domain* a, Domain* b) {
      return std::tie(only_writable ? a->num_picked_writable : a->num_picked,
                      a->priority) >
          std::tie(only_writable ? b->num_picked_writable : b->num_picked,
                   b->priority);
    };
    std::priority_queue<Domain*, std::vector<Domain*>, decltype(cmp)> queue(
        cmp);

    // Initialize queue.
    size_t result_size = 0;
    for (auto& kv : domains) {
      Domain* d = &kv.second;
      result_size += only_writable ? d->num_picked_writable : d->num_picked;
      if (!d->nodes.empty()) {
        queue.push(d);
      }
    }

    // Pick the nodes.
    while (!queue.empty()) {
      Domain* d = queue.top();
      queue.pop();
      assert(!d->nodes.empty());
      int picked_in_domain =
          only_writable ? d->num_picked_writable : d->num_picked;
      if (picked_in_domain >= min_nodes_per_domain &&
          result_size >= target_nodeset_size) {
        // The nodeset is big enough and has enough nodes from each domain.
        break;
      }
      if (res.storage_set.size() >= NODESET_SIZE_MAX) {
        // Nodeset is not allowed to grow any bigger.
        RATELIMIT_INFO(std::chrono::seconds(10),
                       2,
                       "Hit max allowed nodeset size %d when selecting nodeset "
                       "for log %lu. Truncating the nodeset.",
                       (int)NODESET_SIZE_MAX,
                       log_id.val());
        break;
      }
      if (!only_writable && res.storage_set.size() >= NODESET_SIZE_MAX / 2) {
        // We're in the first stage but are already getting close to the max
        // allowed nodeset size. Skip to the second stage early to make sure
        // we pick enough writable nodes.
        break;
      }
      if (only_writable) {
        // Skip unwritable nodes.
        while (!d->nodes.empty() && !d->nodes.back().is_writable) {
          d->nodes.pop_back();
        }
        if (d->nodes.empty()) {
          continue;
        }
      }
      ++d->num_picked;
      if (d->nodes.back().is_writable) {
        ++d->num_picked_writable;
      }
      res.storage_set.push_back(d->nodes.back().shard_id);
      d->nodes.pop_back();
      ++result_size;

      if (!d->nodes.empty()) {
        queue.push(d);
      }
    }
  };

  select_nodes(false);
  select_nodes(true);

  ld_check(res.storage_set.size() <= NODESET_SIZE_MAX);

  // Sort the output.
  std::sort(res.storage_set.begin(), res.storage_set.end());

  // Check that we have enough nodes and domains to satisfy the replication
  // requirement.
  if (!configuration::nodes::validStorageSet(
          nodes_configuration, res.storage_set, replication_property)) {
    RATELIMIT_ERROR(std::chrono::seconds(1),
                    5,
                    "Not enough storage nodes to select nodeset for log %lu, "
                    "replication: %s, selected: %s.",
                    log_id.val_,
                    replication_property.toString().c_str(),
                    toString(res.storage_set).c_str());
    res.decision = Decision::FAILED;
    return res;
  }

  res.decision = (prev && prev->shards == res.storage_set)
      ? Decision::KEEP
      : Decision::NEEDS_CHANGE;
  return res;
}

}} // namespace facebook::logdevice
