Optimize RendezvousHashFunc::operator() to avoid heap construction

Reviewed By: stuclar

Differential Revision: D95628434

fbshipit-source-id: a193cd061d49339ff7b1e090f2b7390ff7e0f67f

Author: Brian Ecker

mcrouter/lib/RendezvousHashFunc.cpp

@@ -44,6 +44,27 @@ std::vector<RendezvousIterator::ScoreAndIndex> get_scores(
 }
 } // namespace
 
+// O(N) max-tracking loop instead of building a full heap via begin().
+// Similar structure to WeightedRendezvousHashFunc::operator(), but uses the
+// raw hash value as the score rather than the weighted inverse-log transform
+// (weight * 1/(-log(U))). Both select the same winner when weights are uniform,
+// but this avoids the log/division cost. Tie-breaking (higher index wins) is
+// kept to match the priority_queue-based iterator path in begin().
+size_t RendezvousHashFunc::operator()(folly::StringPiece key) const {
+  const uint64_t keyHash = RendezvousIterator::keyHash(key);
+  double maxScore = -1.0;
+  size_t maxIndex = 0;
+  for (size_t i = 0; i < endpointHashes_.size(); ++i) {
+    double score =
+        static_cast<double>(hash128to64(endpointHashes_[i], keyHash));
+    if (score > maxScore || (score == maxScore && i > maxIndex)) {
+      maxScore = score;
+      maxIndex = i;
+    }
+  }
+  return maxIndex;
+}
+
 RendezvousIterator RendezvousHashFunc::begin(folly::StringPiece key) const {
   return RendezvousIterator(get_scores(endpointHashes_, key));
 }

mcrouter/lib/RendezvousHashFunc.h

@@ -37,9 +37,7 @@ class RendezvousHashFunc {
       const std::vector<folly::StringPiece>& endpoints,
       const folly::dynamic& json);
 
-  size_t operator()(folly::StringPiece key) const {
-    return *begin(key);
-  }
+  size_t operator()(folly::StringPiece key) const;
 
   RendezvousIterator begin(folly::StringPiece key) const;
 

mcrouter/lib/test/RendezvousHashBenchmark.cpp

@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <folly/Benchmark.h>
+#include <folly/Range.h>
+#include <folly/init/Init.h>
+#include <folly/json/dynamic.h>
+
+#include "mcrouter/lib/RendezvousHashFunc.h"
+
+constexpr folly::StringPiece kKey =
+    "someKey_ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789+_sdkfjsdfksjdfasdfaksxxx";
+
+using facebook::memcache::RendezvousHashFunc;
+
+namespace {
+
+std::vector<folly::StringPiece> makeEndpoints(size_t n) {
+  static std::vector<std::string> storage;
+  storage.clear();
+  storage.reserve(n);
+  for (size_t i = 0; i < n; ++i) {
+    storage.push_back("endpoint" + std::to_string(i) + ".example.com:1234");
+  }
+  std::vector<folly::StringPiece> endpoints;
+  endpoints.reserve(n);
+  for (const auto& s : storage) {
+    endpoints.emplace_back(s);
+  }
+  return endpoints;
+}
+
+} // namespace
+
+void rendezvousHashBench(size_t iters, size_t poolSize) {
+  RendezvousHashFunc func(
+      std::vector<folly::StringPiece>(), folly::dynamic::object);
+  BENCHMARK_SUSPEND {
+    auto endpoints = makeEndpoints(poolSize);
+    func = RendezvousHashFunc(endpoints, folly::dynamic::object);
+  }
+  for (size_t i = 0; i < iters; ++i) {
+    folly::doNotOptimizeAway(func(kKey));
+  }
+}
+
+BENCHMARK_NAMED_PARAM(rendezvousHashBench, pool_10, 10)
+BENCHMARK_NAMED_PARAM(rendezvousHashBench, pool_100, 100)
+BENCHMARK_NAMED_PARAM(rendezvousHashBench, pool_500, 500)
+BENCHMARK_NAMED_PARAM(rendezvousHashBench, pool_1000, 1000)
+
+int main(int argc, char** argv) {
+  const folly::Init init(&argc, &argv, true);
+
+  folly::runBenchmarks();
+  return 0;
+}

mcrouter/lib/test/RendezvousHashTest.cpp

@@ -73,6 +73,14 @@ TEST(RendezvousHashFunc, rendezvous_10) {
           {947, 1026, 1028, 981, 1016, 970, 1013, 939, 1023, 1057}));
 }
 
+TEST(RendezvousHashFunc, operator_matches_begin) {
+  auto func = genRendezvousHashFunc(343);
+  for (size_t i = 0; i < 10000; ++i) {
+    auto key = "mykey:" + folly::to<std::string>(i);
+    EXPECT_EQ(func(key), *func.begin(key));
+  }
+}
+
 TEST(RendezvousHashFunc, rendezvous_rehash) {
   const uint32_t n = 499;
   auto combined = test::genEndpoints(n);