Add support to split keys with hash stop

Summary: Currently hash stops do no support key splitting, this diff adds support. Reviewed By: alikhtarov Differential Revision: D53367269 fbshipit-source-id: c123e14a757f9ccbb8bf39a84dcb23fb4b69367b
facebook · Feb 9, 2024 · fb23451 · fb23451
1 parent 202794e
commit fb23451
Show file tree

Hide file tree

Showing 5 changed files with 242 additions and 19 deletions.
diff --git a/mcrouter/lib/carbon/Keys-inl.h b/mcrouter/lib/carbon/Keys-inl.h
@@ -71,7 +71,8 @@ void Keys<Storage>::update() {
   if (pos != std::string::npos) {
     pos = keyWithoutRoute_.find("|#|", pos);
     if (pos != std::string::npos) {
-      routingKey_.reset(keyWithoutRoute_.begin(), pos);
+      routingKey_ = keyWithoutRoute_.subpiece(0, pos);
+      afterRoutingKey_ = keyWithoutRoute_.subpiece(pos);
     }
   }
   routingKeyHash_ = 0;

diff --git a/mcrouter/lib/carbon/Keys.h b/mcrouter/lib/carbon/Keys.h
@@ -41,6 +41,7 @@ inline folly::IOBuf makeKey<folly::IOBuf>(folly::StringPiece sp) {
  * keyWithoutRoute:                       ^^^^^^^^^^^^^
  * routingPrefix:         ^^^^^^^^^^^^^^^^
  * routingKey:                            ^^^^^^^
+ * afterRoutingKey:                              ^^^^^^
  */
 template <class Storage>
 class Keys {
@@ -109,6 +110,11 @@ class Keys {
   folly::StringPiece routingKey() const {
     return routingKey_;
   }
+
+  folly::StringPiece afterRoutingKey() const {
+    return afterRoutingKey_;
+  }
+
   uint32_t routingKeyHash() const {
     if (!routingKeyHash_) {
       const auto keyPiece = routingKey();
@@ -181,6 +187,7 @@ class Keys {
     keyWithoutRoute_ = other.keyWithoutRoute_;
     routingPrefix_ = other.routingPrefix_;
     routingKey_ = other.routingKey_;
+    afterRoutingKey_ = other.afterRoutingKey_;
     routingKeyHash_ = other.routingKeyHash_;
     lastHash_.size_ = other.lastHash_.size_;
     lastHash_.hash_ = other.lastHash_.hash_;
@@ -212,6 +219,7 @@ class Keys {
   folly::StringPiece keyWithoutRoute_;
   folly::StringPiece routingPrefix_;
   folly::StringPiece routingKey_;
+  folly::StringPiece afterRoutingKey_;
   mutable uint32_t routingKeyHash_{0};
 
   struct HashData {

diff --git a/mcrouter/lib/carbon/test/CarbonStructuresTest.cpp b/mcrouter/lib/carbon/test/CarbonStructuresTest.cpp
@@ -55,9 +55,11 @@ void checkKeyFilledProperly(const Key& key) {
   EXPECT_EQ(std::strlen(kKeyLiteral), key.size());
   EXPECT_EQ(kKeyLiteral, key.fullKey());
   EXPECT_EQ(
-      "abcdefghijklmnopqrstuvwxyz|#|afterhashstop", key.keyWithoutRoute());
-  EXPECT_EQ("/region/cluster/", key.routingPrefix());
-  EXPECT_EQ("abcdefghijklmnopqrstuvwxyz", key.routingKey());
+      "abcdefghijklmnopqrstuvwxyz|#|afterhashstop",
+      key.keyWithoutRoute().str());
+  EXPECT_EQ("/region/cluster/", key.routingPrefix().str());
+  EXPECT_EQ("abcdefghijklmnopqrstuvwxyz", key.routingKey().str());
+  EXPECT_EQ("|#|afterhashstop", key.afterRoutingKey().str());
   EXPECT_NE(0, key.routingKeyHash());
   EXPECT_TRUE(key.hasHashStop());
 }
@@ -261,6 +263,7 @@ TEST(CarbonBasic, setAndGet) {
   EXPECT_EQ(
       "abcdefghijklmnopqrstuvwxyz|#|afterhashstop",
       req.key_ref()->keyWithoutRoute().str());
+  EXPECT_EQ("|#|afterhashstop", req.key_ref()->afterRoutingKey().str());
 
   const auto reqKeyPiece2 = req2.key_ref()->fullKey();
   EXPECT_EQ(kKeyLiteral, reqKeyPiece2);
@@ -270,6 +273,7 @@ TEST(CarbonBasic, setAndGet) {
   EXPECT_EQ(
       "abcdefghijklmnopqrstuvwxyz|#|afterhashstop",
       req2.key_ref()->keyWithoutRoute().str());
+  EXPECT_EQ("|#|afterhashstop", req2.key_ref()->afterRoutingKey().str());
 
   // bool
   req.testBool_ref() = true;
@@ -634,6 +638,7 @@ TEST(CarbonBasic, setAndGetFieldRefAPI) {
   EXPECT_EQ(
       "abcdefghijklmnopqrstuvwxyz|#|afterhashstop",
       req.key_ref()->keyWithoutRoute().str());
+  EXPECT_EQ("|#|afterhashstop", req.key_ref()->afterRoutingKey().str());
 
   const auto reqKeyPiece2 = req2.key_ref()->fullKey();
   EXPECT_EQ(kKeyLiteral, reqKeyPiece2);
@@ -643,6 +648,7 @@ TEST(CarbonBasic, setAndGetFieldRefAPI) {
   EXPECT_EQ(
       "abcdefghijklmnopqrstuvwxyz|#|afterhashstop",
       req2.key_ref()->keyWithoutRoute().str());
+  EXPECT_EQ("|#|afterhashstop", req2.key_ref()->afterRoutingKey().str());
 
   // bool
   req.testBool_ref() = true;
@@ -760,6 +766,7 @@ TEST(CarbonBasic, setAndGetFieldRefAPIThrift) {
   EXPECT_EQ(
       "abcdefghijklmnopqrstuvwxyz|#|afterhashstop",
       req.key_ref()->keyWithoutRoute().str());
+  EXPECT_EQ("|#|afterhashstop", req.key_ref()->afterRoutingKey().str());
 
   // bool
   req.testBool_ref() = true;

diff --git a/mcrouter/routes/KeySplitRoute.h b/mcrouter/routes/KeySplitRoute.h
@@ -185,8 +185,16 @@ class KeySplitRoute {
   template <class Request>
   Request copyAndAugment(Request& originalReq, uint64_t replicaId) const {
     auto req = originalReq;
-    req.key_ref() = folly::to<std::string>(
-        req.key_ref()->fullKey(), kMemcacheReplicaSeparator, replicaId);
+    if (req.key_ref()->hasHashStop()) {
+      req.key_ref() = folly::to<std::string>(
+          req.key_ref()->routingKey(),
+          kMemcacheReplicaSeparator,
+          replicaId,
+          req.key_ref()->afterRoutingKey());
+    } else {
+      req.key_ref() = folly::to<std::string>(
+          req.key_ref()->fullKey(), kMemcacheReplicaSeparator, replicaId);
+    }
     return req;
   }