HashStopAllowListRoute

Summary: Allow list hashstops in production.

Reviewed By: lenar-f

Differential Revision: D54132835

fbshipit-source-id: bd12fdf80dc2280c2dbda8f4b7b58d381b372fe4

mcrouter/routes/HashStopAllowListRoute.cpp

@@ -0,0 +1,35 @@
+/*
+ * 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 "mcrouter/routes/HashStopAllowListRoute.h"
+
+#include <folly/Format.h>
+#include <folly/fibers/WhenN.h>
+
+namespace facebook::memcache::mcrouter {
+
+HashStopAllowListRouteSettings parseHashStopAllowListRouteSettings(
+    const folly::dynamic& json) {
+  HashStopAllowListRouteSettings settings;
+  auto jPrefixList = json.get_ptr("prefixes");
+  if (jPrefixList && jPrefixList->isArray() && jPrefixList->size() > 0) {
+    std::vector<std::pair<std::string, bool>> prefixes;
+    for (size_t i = 0; i < jPrefixList->size(); ++i) {
+      const auto& prefix = jPrefixList->at(i);
+      checkLogic(
+          prefix.isString(),
+          "{} expected string, found {}",
+          prefix,
+          prefix.typeName());
+      prefixes.emplace_back(prefix.asString(), true);
+    }
+    settings.prefixes = std::move(prefixes);
+  }
+  return settings;
+}
+
+} // namespace facebook::memcache::mcrouter

mcrouter/routes/HashStopAllowListRoute.h

@@ -0,0 +1,104 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "mcrouter/McrouterFiberContext.h"
+#include "mcrouter/lib/Ch3HashFunc.h"
+#include "mcrouter/lib/Reply.h"
+#include "mcrouter/lib/RouteHandleTraverser.h"
+#include "mcrouter/lib/config/RouteHandleFactory.h"
+#include "mcrouter/lib/fbi/cpp/LowerBoundPrefixMap.h"
+#include "mcrouter/lib/fbi/cpp/ParsingUtil.h"
+#include "mcrouter/routes/RoutingUtils.h"
+
+namespace facebook::memcache::mcrouter {
+
+struct HashStopAllowListRouteSettings {
+  std::vector<std::pair<std::string, bool>> prefixes;
+};
+
+/**
+ * This route handle allow lists uses of the mcrouter HashStop.
+ * Keys which have a hashstop and are not configured in the allow
+ * list will return LOCAL_ERROR.
+ */
+template <class RouterInfo>
+class HashStopAllowListRoute {
+ private:
+  using RouteHandleIf = typename RouterInfo::RouteHandleIf;
+  using RouteHandlePtr = typename RouterInfo::RouteHandlePtr;
+
+ public:
+  HashStopAllowListRoute(
+      RouteHandlePtr rh,
+      HashStopAllowListRouteSettings& settings)
+      : rh_(std::move(rh)), prefixMap_(settings.prefixes) {}
+
+  std::string routeName() const {
+    return fmt::format(
+        "HashStopAllowListRoute|num_prefixes={}", prefixMap_.size());
+  }
+
+  template <class Request>
+  bool traverse(
+      const Request& req,
+      const RouteHandleTraverser<RouteHandleIf>& t) const {
+    if (FOLLY_LIKELY(!req.key_ref()->hasHashStop()) ||
+        allowedHashStop(req.key_ref()->fullKey())) {
+      return t(*rh_, req);
+    }
+    return false;
+  }
+
+  template <class Request>
+  ReplyT<Request> route(const Request& req) const {
+    if (FOLLY_LIKELY(!req.key_ref()->hasHashStop()) ||
+        allowedHashStop(req.key_ref()->fullKey())) {
+      return rh_->route(req);
+    }
+    return createReply<Request>(
+        ErrorReply, carbon::Result::LOCAL_ERROR, "Hash stop not permited");
+  }
+
+ private:
+  const RouteHandlePtr rh_;
+  const memcache::LowerBoundPrefixMap<bool> prefixMap_;
+
+  FOLLY_ALWAYS_INLINE bool allowedHashStop(const folly::StringPiece key) const {
+    if (prefixMap_.empty()) {
+      return false;
+    }
+    auto it = prefixMap_.findPrefix(key);
+    return it != prefixMap_.end();
+  }
+};
+
+HashStopAllowListRouteSettings parseHashStopAllowListRouteSettings(
+    const folly::dynamic& json);
+
+template <class RouterInfo>
+typename RouterInfo::RouteHandlePtr createHashStopAllowListRoute(
+    typename RouterInfo::RouteHandlePtr rh,
+    const folly::dynamic& json) {
+  auto settings = parseHashStopAllowListRouteSettings(json);
+  return makeRouteHandleWithInfo<RouterInfo, HashStopAllowListRoute>(
+      std::move(rh), settings);
+}
+
+template <class RouterInfo>
+typename RouterInfo::RouteHandlePtr makeHashStopAllowListRoute(
+    RouteHandleFactory<typename RouterInfo::RouteHandleIf>& factory,
+    const folly::dynamic& json) {
+  checkLogic(json.isObject(), "HashStopAllowListRoute is not an object");
+  checkLogic(
+      json.count("target"), "HashStopAllowListRoute: Missing target parameter");
+  return createHashStopAllowListRoute<RouterInfo>(
+      factory.create(json["target"]), json);
+}
+
+} // namespace facebook::memcache::mcrouter

mcrouter/routes/McRouteHandleProvider.cpp

@@ -23,6 +23,7 @@
 #include "mcrouter/routes/FailoverRoute.h"
 #include "mcrouter/routes/FailoverWithExptimeRouteFactory.h"
 #include "mcrouter/routes/HashRouteFactory.h"
+#include "mcrouter/routes/HashStopAllowListRoute.h"
 #include "mcrouter/routes/HostIdRouteFactory.h"
 #include "mcrouter/routes/KeyParseRoute.h"
 #include "mcrouter/routes/KeySplitRoute.h"
@@ -268,6 +269,8 @@ McRouteHandleProvider<MemcacheRouterInfo>::buildRouteMap() {
        [](McRouteHandleFactory& factory, const folly::dynamic& json) {
          return makeHashRoute<McrouterRouterInfo>(factory, json);
        }},
+      {"HashStopAllowListRoute",
+       &makeHashStopAllowListRoute<MemcacheRouterInfo>},
       {"HostIdRoute", &makeHostIdRoute<MemcacheRouterInfo>},
       {"LatencyInjectionRoute", &makeLatencyInjectionRoute<MemcacheRouterInfo>},
       {"L1L2CacheRoute", &makeL1L2CacheRoute<MemcacheRouterInfo>},

mcrouter/routes/test/HashStopAllowListRouteTest.cpp

@@ -0,0 +1,169 @@
+/*
+ * 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 <algorithm>
+#include <memory>
+#include <vector>
+
+#include <gtest/gtest.h>
+
+#include "mcrouter/lib/carbon/example/gen/HelloGoodbyeRouterInfo.h"
+#include "mcrouter/lib/test/RouteHandleTestUtil.h"
+#include "mcrouter/lib/test/TestRouteHandle.h"
+
+#include "mcrouter/routes/HashStopAllowListRoute.h"
+
+#include "mcrouter/routes/test/RouteHandleTestBase.h"
+#include "mcrouter/routes/test/RouteHandleTestUtil.h"
+
+using namespace hellogoodbye;
+using namespace facebook::memcache;
+using namespace facebook::memcache::mcrouter;
+
+using std::make_shared;
+using std::vector;
+
+namespace facebook::memcache::mcrouter {
+
+using TestHandle = TestHandleImpl<MemcacheRouteHandleIf>;
+using TestRouteHandle = std::shared_ptr<MemcacheRouteHandleIf>;
+using HashStopAllowListRouteHandle =
+    McrouterRouteHandle<HashStopAllowListRoute<MemcacheRouterInfo>>;
+
+std::pair<std::vector<std::shared_ptr<TestHandle>>, TestRouteHandle>
+getMockRouteHandle() {
+  std::vector<std::shared_ptr<TestHandle>> handleVec{
+      std::make_shared<TestHandle>(
+          GetRouteTestData(carbon::Result::FOUND, "a")),
+  };
+  return {std::move(handleVec), get_route_handles(handleVec)[0]};
+}
+
+TEST(HashStopAllowListRouteTest, NoMatch) {
+  constexpr folly::StringPiece kRoutingConfig = R"(
+  {
+    "prefixes": [
+        "foo",
+        "baz"
+    ]
+  }
+  )";
+
+  auto [handleVec, mockRouteHandle] = getMockRouteHandle();
+  auto rh = createHashStopAllowListRoute<MemcacheRouterInfo>(
+      mockRouteHandle, folly::parseJson(kRoutingConfig));
+  std::string key = "abc";
+
+  McSetRequest reqSet(key);
+  reqSet.value_ref() = folly::IOBuf(folly::IOBuf::COPY_BUFFER, "value");
+
+  auto reply = rh->route(reqSet);
+
+  EXPECT_FALSE(handleVec[0]->saw_keys.empty());
+  EXPECT_EQ(key, handleVec[0]->saw_keys[0]);
+}
+
+TEST(HashStopAllowListRouteTest, HashStopRejected) {
+  constexpr folly::StringPiece kRoutingConfig = R"(
+  {
+    "prefixes": [
+        "foo",
+        "baz"
+    ]
+  }
+  )";
+
+  auto [handleVec, mockRouteHandle] = getMockRouteHandle();
+  auto rh = createHashStopAllowListRoute<MemcacheRouterInfo>(
+      mockRouteHandle, folly::parseJson(kRoutingConfig));
+  std::string key = "abc|#|1";
+
+  McSetRequest reqSet(key);
+  reqSet.value_ref() = folly::IOBuf(folly::IOBuf::COPY_BUFFER, "value");
+
+  auto reply = rh->route(reqSet);
+
+  EXPECT_EQ(carbon::Result::LOCAL_ERROR, *reply.result_ref());
+  EXPECT_TRUE(handleVec[0]->saw_keys.empty());
+}
+
+TEST(HashStopAllowListRouteTest, HashStopAccepted) {
+  constexpr folly::StringPiece kRoutingConfig = R"(
+  {
+    "prefixes": [
+        "foo",
+        "baz",
+        "abc"
+    ]
+  }
+  )";
+
+  auto [handleVec, mockRouteHandle] = getMockRouteHandle();
+  auto rh = createHashStopAllowListRoute<MemcacheRouterInfo>(
+      mockRouteHandle, folly::parseJson(kRoutingConfig));
+  std::string key = "abc|#|1";
+
+  McSetRequest reqSet(key);
+  reqSet.value_ref() = folly::IOBuf(folly::IOBuf::COPY_BUFFER, "value");
+
+  auto reply = rh->route(reqSet);
+
+  EXPECT_FALSE(handleVec[0]->saw_keys.empty());
+  EXPECT_EQ(key, handleVec[0]->saw_keys[0]);
+}
+
+TEST(HashStopAllowListRouteTest, HashStopAcceptAndReject) {
+  constexpr folly::StringPiece kRoutingConfig = R"(
+  {
+    "prefixes": [
+        "foo",
+        "baz",
+        "abc"
+    ]
+  }
+  )";
+
+  auto [handleVec, mockRouteHandle] = getMockRouteHandle();
+  auto rh = createHashStopAllowListRoute<MemcacheRouterInfo>(
+      mockRouteHandle, folly::parseJson(kRoutingConfig));
+  std::string key = "zoo";
+
+  McSetRequest reqSet(key);
+  reqSet.value_ref() = folly::IOBuf(folly::IOBuf::COPY_BUFFER, "ooz");
+  auto reply = rh->route(reqSet);
+  EXPECT_FALSE(handleVec[0]->saw_keys.empty());
+  EXPECT_EQ(key, handleVec[0]->saw_keys[0]);
+
+  key = "zoo|#|a";
+  reqSet = McSetRequest(key);
+  reqSet.value_ref() = folly::IOBuf(folly::IOBuf::COPY_BUFFER, "b");
+  reply = rh->route(reqSet);
+  EXPECT_EQ(carbon::Result::LOCAL_ERROR, *reply.result_ref());
+
+  key = "foo|#|a";
+  reqSet = McSetRequest(key);
+  reqSet.value_ref() = folly::IOBuf(folly::IOBuf::COPY_BUFFER, "b");
+  reply = rh->route(reqSet);
+  EXPECT_FALSE(handleVec[0]->saw_keys.empty());
+  EXPECT_EQ(key, handleVec[0]->saw_keys[1]);
+
+  key = "baz|#|a";
+  reqSet = McSetRequest(key);
+  reqSet.value_ref() = folly::IOBuf(folly::IOBuf::COPY_BUFFER, "b");
+  reply = rh->route(reqSet);
+  EXPECT_FALSE(handleVec[0]->saw_keys.empty());
+  EXPECT_EQ(key, handleVec[0]->saw_keys[2]);
+
+  key = "abc|#|a";
+  reqSet = McSetRequest(key);
+  reqSet.value_ref() = folly::IOBuf(folly::IOBuf::COPY_BUFFER, "b");
+  reply = rh->route(reqSet);
+  EXPECT_FALSE(handleVec[0]->saw_keys.empty());
+  EXPECT_EQ(key, handleVec[0]->saw_keys[3]);
+}
+
+} // namespace facebook::memcache::mcrouter