Send request-key dedup key to Axon DL write proxy

Summary:
The Axon DL write proxy (dl/write_proxy/service/WriteProxyServer.cpp)
supports request dedup keyed on cluster + keyspace + DLAppendRequest.dedupeKey,
but mcrouter never set dedupeKey, so the dedup window was effectively unusable
from the Memcache invalidation path.

This sets dedupeKey in AxonLogRoute::writeProxy to the request key
(req.key()->fullKey()) for both sets and deletes. All writes to the same key
within the proxy dedup window therefore collapse, which is what suppresses hot
keys (the goal of this feature). Per review discussion, writes to the same key
are treated as concurrent, so collapsing them to a single append is acceptable.

The proxy dedup is best-effort, in-memory and per-host, and is gated off by
default (append_dedup_window_sec=0). A follow-up will set a small dedup window
(<=100ms) once the feature is enabled.

Server-side consumption already exists in the DL write proxy; no change needed
there.

Reviewed By: lenar-f

Differential Revision: D107702706

fbshipit-source-id: 3a77b58240a9fc36199ab4a12730b37e7648dcdf

Author: Keshav Sharma

mcrouter/McDistributionUtils.cpp

@@ -8,6 +8,8 @@
 #include "mcrouter/McDistributionUtils.h"
 
 #include <optional>
+#include <string>
+#include <utility>
 
 namespace facebook {
 namespace memcache {
@@ -35,6 +37,8 @@ FOLLY_NOINLINE folly::exception_wrapper distributeWriteRequest(
     pool.emplace(axonCtx->poolFilter);
   }
 
+  // Dedupe by request key to suppress hot keys at the write proxy.
+  auto dedupeKey = req.key()->fullKey().str();
   // Run serialization off fiber
   auto kvPairs = folly::fibers::runInMainContext([&req,
                                                   &targetRegion,
@@ -57,7 +61,8 @@ FOLLY_NOINLINE folly::exception_wrapper distributeWriteRequest(
         invalidation::DistributionOperation::Write,
         sourceRegion);
   });
-  return axonCtx->writeProxyFn(bucketId, std::move(kvPairs), secureWrites);
+  return axonCtx->writeProxyFn(
+      bucketId, std::move(kvPairs), secureWrites, std::move(dedupeKey));
 }
 
 FOLLY_NOINLINE folly::exception_wrapper distributeDeleteRequest(
@@ -79,6 +84,8 @@ FOLLY_NOINLINE folly::exception_wrapper distributeDeleteRequest(
   if (!axonCtx->poolFilter.empty()) {
     pool.emplace(axonCtx->poolFilter);
   }
+  // Dedupe by request key to suppress hot keys at the write proxy.
+  auto dedupeKey = req.key()->fullKey().str();
   // Run off fiber to save fiber stack for serialization
   auto kvPairs = folly::fibers::runInMainContext(
       [&req, &region, &pool, &message, &sourceRegion, type]() {
@@ -110,7 +117,10 @@ FOLLY_NOINLINE folly::exception_wrapper distributeDeleteRequest(
             sourceRegion);
       });
   return axonCtx->writeProxyFn(
-      bucketId, std::move(kvPairs), /*secureWrites*/ false);
+      bucketId,
+      std::move(kvPairs),
+      /*secureWrites*/ false,
+      std::move(dedupeKey));
 }
 
 FOLLY_NOINLINE bool spoolAsynclog(

mcrouter/McrouterFiberContext.h

@@ -10,6 +10,8 @@
 #include <bitset>
 #include <map>
 #include <memory>
+#include <optional>
+#include <string>
 #include <utility>
 
 #include <folly/Range.h>
@@ -60,7 +62,8 @@ using ExtraDataCallbackT = std::function<ExtraDataMap()>;
 using AxonProxyWriteFn = std::function<folly::exception_wrapper(
     uint64_t,
     folly::F14FastMap<std::string, std::string>&&,
-    bool)>;
+    bool,
+    std::optional<std::string> /* dedupeKey */)>;
 
 struct AxonContext {
   bool fallbackAsynclog{false};

mcrouter/routes/test/DistributionRouteTest.cpp

@@ -34,8 +34,10 @@ struct AxonMockResult {
 namespace {
 AxonMockResult setupAxonFn(std::shared_ptr<AxonContext>& ctx) {
   AxonMockResult res;
-  ctx->writeProxyFn =
-      [&](auto bucketId, auto&& payload, bool) -> folly::exception_wrapper {
+  ctx->writeProxyFn = [&](auto bucketId,
+                          auto&& payload,
+                          bool,
+                          auto&& /* dedupeKey */) -> folly::exception_wrapper {
     res.bucketId = bucketId;
     if (payload.find(invalidation::kRegion) != payload.end()) {
       res.region = payload.at(invalidation::kRegion);

mcrouter/test/cpp_unit_tests/McDistributionUtilsTest.cpp

@@ -26,7 +26,8 @@ TEST(McDistributionUtilsTest, distributeSetTest) {
   facebook::memcache::mcrouter::AxonProxyWriteFn writeProxyFn =
       [&](auto bucketId,
           auto&& kvPairs,
-          bool isSecureWrites) -> folly::exception_wrapper {
+          bool isSecureWrites,
+          const std::optional<std::string>&) -> folly::exception_wrapper {
     resultBucketId = bucketId;
     resultKvPairs = kvPairs;
     secureWrites = isSecureWrites;
@@ -96,7 +97,8 @@ TEST(McDistributionUtilsTest, distributeSetWithSecurityTest) {
   facebook::memcache::mcrouter::AxonProxyWriteFn writeProxyFn =
       [&](auto bucketId,
           auto&& kvPairs,
-          bool isSecureWrites) -> folly::exception_wrapper {
+          bool isSecureWrites,
+          const std::optional<std::string>&) -> folly::exception_wrapper {
     resultBucketId = bucketId;
     resultKvPairs = kvPairs;
     secureWrites = isSecureWrites;
@@ -169,7 +171,10 @@ TEST(McDistributionUtilsTest, distributeDeleteDirectedTest) {
   folly::F14FastMap<std::string, std::string> resultKvPairs;
 
   facebook::memcache::mcrouter::AxonProxyWriteFn writeProxyFn =
-      [&](auto bucketId, auto&& kvPairs, bool) -> folly::exception_wrapper {
+      [&](auto bucketId,
+          auto&& kvPairs,
+          bool,
+          const std::optional<std::string>&) -> folly::exception_wrapper {
     resultBucketId = bucketId;
     resultKvPairs = kvPairs;
     return folly::exception_wrapper();
@@ -242,7 +247,10 @@ TEST(McDistributionUtilsTest, distributeDeleteBroadcastTest) {
   folly::F14FastMap<std::string, std::string> resultKvPairs;
 
   facebook::memcache::mcrouter::AxonProxyWriteFn writeProxyFn =
-      [&](auto bucketId, auto&& kvPairs, bool) -> folly::exception_wrapper {
+      [&](auto bucketId,
+          auto&& kvPairs,
+          bool,
+          const std::optional<std::string>&) -> folly::exception_wrapper {
     resultBucketId = bucketId;
     resultKvPairs = kvPairs;
     return folly::exception_wrapper();
@@ -314,7 +322,10 @@ TEST(McDistributionUtilsTest, spoolDeleteTest) {
   folly::F14FastMap<std::string, std::string> resultKvPairs;
 
   facebook::memcache::mcrouter::AxonProxyWriteFn writeProxyFn =
-      [&](auto bucketId, auto&& kvPairs, bool) -> folly::exception_wrapper {
+      [&](auto bucketId,
+          auto&& kvPairs,
+          bool,
+          const std::optional<std::string>&) -> folly::exception_wrapper {
     resultBucketId = bucketId;
     resultKvPairs = kvPairs;
     return folly::exception_wrapper();
@@ -376,4 +387,67 @@ TEST(McDistributionUtilsTest, spoolDeleteTest) {
       DistributionOperation::Delete);
   EXPECT_EQ(resultKvPairs.find(kSourceRegion)->second, sourceRegion);
 }
+
+// The dedup key passed to the write proxy is the request key, so all writes to
+// the same key collapse within the proxy's dedup window (hot-key suppression)
+// regardless of value, while writes to different keys never collapse.
+TEST(McDistributionUtilsTest, distributeSetKeyBasedDedupeKeyTest) {
+  std::string targetRegion = "texas";
+  std::string sourceRegion = "altoona";
+  std::string pool = "main";
+  uint64_t bucketId = 234532;
+  std::optional<std::string> resultDedupeKey;
+
+  facebook::memcache::mcrouter::AxonProxyWriteFn writeProxyFn =
+      [&](auto,
+          auto&&,
+          bool,
+          std::optional<std::string> dedupeKey) -> folly::exception_wrapper {
+    resultDedupeKey = std::move(dedupeKey);
+    return folly::exception_wrapper();
+  };
+
+  auto axonCtx = mcrouter::AxonContext{
+      .fallbackAsynclog = false,
+      .allDelete = false,
+      .writeProxyFn = writeProxyFn,
+      .defaultRegionFilter = sourceRegion,
+      .poolFilter = pool};
+  auto axonCtxPtr =
+      std::make_shared<facebook::memcache::mcrouter::AxonContext>(axonCtx);
+
+  auto makeReq = [](folly::StringPiece key, folly::StringPiece value) {
+    auto req = McSetRequest(key);
+    req.value() = folly::IOBuf(folly::IOBuf::COPY_BUFFER, value);
+    return req;
+  };
+
+  // Dedup key is the request key itself.
+  auto req = makeReq("key1", "value");
+  EXPECT_FALSE(distributeWriteRequest(
+                   req, axonCtxPtr, bucketId, targetRegion, sourceRegion)
+                   .has_exception_ptr());
+  ASSERT_TRUE(resultDedupeKey.has_value());
+  EXPECT_EQ(*resultDedupeKey, "key1");
+  const auto firstKey = *resultDedupeKey;
+
+  // Same key, different value -> same dedup key (collapses the hot key).
+  resultDedupeKey.reset();
+  auto sameKeyReq = makeReq("key1", "value2");
+  EXPECT_FALSE(distributeWriteRequest(
+                   sameKeyReq, axonCtxPtr, bucketId, targetRegion, sourceRegion)
+                   .has_exception_ptr());
+  ASSERT_TRUE(resultDedupeKey.has_value());
+  EXPECT_EQ(*resultDedupeKey, firstKey);
+
+  // Different key -> different dedup key (must NOT collapse).
+  resultDedupeKey.reset();
+  auto otherKeyReq = makeReq("key2", "value");
+  EXPECT_FALSE(
+      distributeWriteRequest(
+          otherKeyReq, axonCtxPtr, bucketId, targetRegion, sourceRegion)
+          .has_exception_ptr());
+  ASSERT_TRUE(resultDedupeKey.has_value());
+  EXPECT_NE(*resultDedupeKey, firstKey);
+}
 } // namespace facebook::memcache::invalidation::test