Update injection service to populate usecaseId

Summary: As described in D62582055, update injection service to set usecase ID in delete request for DL. D60796664 parse usecase ID for strict namespace DT. The diff invoke the same parsing function. Usecase ID can also be set from PrefixUsecaseRoute in Mcrouter, the diff add support to obtain from recording context. Reviewed By: stuclar Differential Revision: D62582054 fbshipit-source-id: 6b083e5be7152b1fbaf40b279eeb44ac8044097a
facebook · Sep 16, 2024 · 0b2cd13 · 0b2cd13
1 parent 9beaf18
commit 0b2cd13
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 14 deletions.
diff --git a/mcrouter/ProxyRequestContext.h b/mcrouter/ProxyRequestContext.h
@@ -50,8 +50,11 @@ class ProxyRequestContext {
   using ClientCallback =
       std::function<void(const PoolContext&, const AccessPoint&)>;
   using ShardSplitCallback = std::function<void(const ShardSplitter&, bool)>;
-  using BucketizationCallback =
-      std::function<void(std::string, const uint64_t, const std::string&)>;
+  using BucketizationCallback = std::function<void(
+      std::string,
+      const uint64_t,
+      const std::string&,
+      std::optional<int64_t>)>;
 
   virtual ~ProxyRequestContext();
 
@@ -80,13 +83,15 @@ class ProxyRequestContext {
     }
   }
 
+  template <class Request>
   void recordBucketizationData(
       std::string key,
       const uint64_t bucketId,
-      const std::string& keyspace) const {
+      const std::string& keyspace,
+      const Request& req) const {
     if (recording_ && recordingState_->bucketizationCallback) {
       recordingState_->bucketizationCallback(
-          std::move(key), bucketId, keyspace);
+          std::move(key), bucketId, keyspace, getUsecaseId(req));
     }
   }
 

diff --git a/mcrouter/ServiceInfo-inl.h b/mcrouter/ServiceInfo-inl.h
@@ -221,20 +221,22 @@ class GetBucketCommandDispatcher {
       folly::StringPiece keyStr,
       Proxy<RouterInfo>& proxy,
       const ProxyRoute<RouterInfo>& proxyRoute) {
-    using KeyBucketKeyspaceTuple =
-        std::tuple<std::string, std::string, std::string>;
+    using KeyBucketKeyspaceTuple = std::
+        tuple<std::string, std::string, std::string, std::optional<int64_t>>;
     proxy.fiberManager().addTaskFinally(
         [keyStr, &proxy, &proxyRoute]() {
           auto keyBucketKeyspaceTuples =
               std::make_unique<std::vector<KeyBucketKeyspaceTuple>>();
           auto cb = [&keyBucketKeyspaceTuples](
                         std::string keyRecorded,
                         const uint64_t bucketIdRecorded,
-                        const std::string_view bucketizationKeyspaceRecorded) {
+                        const std::string_view bucketizationKeyspaceRecorded,
+                        std::optional<int64_t> usecaseId) {
             keyBucketKeyspaceTuples->push_back(std::make_tuple(
                 std::move(keyRecorded),
                 folly::to<std::string>(bucketIdRecorded),
-                folly::to<std::string>(bucketizationKeyspaceRecorded)));
+                folly::to<std::string>(bucketizationKeyspaceRecorded),
+                std::move(usecaseId)));
           };
           folly::fibers::Baton baton;
           auto rctx =
@@ -254,7 +256,7 @@ class GetBucketCommandDispatcher {
                   data) {
           std::string str;
           const auto& tuplesPtr = *data;
-          for (const auto& [key, bucket, keyspace] : *tuplesPtr) {
+          for (const auto& [key, bucket, keyspace, usecaseId] : *tuplesPtr) {
             if (!str.empty()) {
               str.append("\r\n");
             }
@@ -263,6 +265,10 @@ class GetBucketCommandDispatcher {
             str.append(bucket);
             str.push_back('\t');
             str.append(keyspace);
+            str.push_back('\t');
+            str.append(
+                usecaseId.has_value() ? folly::to<std::string>(*usecaseId)
+                                      : "null");
           }
           ReplyT<ServiceInfoRequest> reply(carbon::Result::FOUND);
           reply.value_ref() = folly::IOBuf(folly::IOBuf::COPY_BUFFER, str);

diff --git a/mcrouter/lib/network/MessageHelpers.h b/mcrouter/lib/network/MessageHelpers.h
@@ -217,8 +217,7 @@ class HasUsecaseIdTrait : public std::false_type {};
 template <typename Message>
 class HasUsecaseIdTrait<
     Message,
-    std::void_t<
-        decltype(std::declval<std::decay_t<Message>>().usecaseId_ref())>>
+    std::void_t<decltype(std::declval<Message>().usecaseId_ref())>>
     : public std::true_type {};
 
 template <typename Message>

diff --git a/mcrouter/routes/McBucketRoute.h b/mcrouter/routes/McBucketRoute.h
@@ -101,7 +101,8 @@ class McBucketRoute {
       ctx->recordBucketizationData(
           req.key_ref()->keyWithoutRoute().str(),
           bucketId,
-          bucketizationKeyspace_);
+          bucketizationKeyspace_,
+          req);
     }
     return fiber_local<RouterInfo>::runWithLocals([this, &req, &t, bucketId]() {
       fiber_local<RouterInfo>::setBucketId(bucketId);
@@ -120,7 +121,8 @@ class McBucketRoute {
       ctx->recordBucketizationData(
           req.key_ref()->keyWithoutRoute().str(),
           bucketId,
-          bucketizationKeyspace_);
+          bucketizationKeyspace_,
+          req);
       return createReply<Request>(DefaultReply, req);
     }
     return routeImpl(req, bucketId);

diff --git a/mcrouter/routes/test/RouteHandleTestUtil.h b/mcrouter/routes/test/RouteHandleTestUtil.h
@@ -65,7 +65,8 @@ void setRecordBucketDataContext(
   auto cb = [&pairs](
                 std::string key,
                 const uint64_t bucket,
-                const std::string_view /*poolRecorded*/) {
+                const std::string_view /* poolRecorded */,
+                std::optional<int64_t> /* usecaseId */) {
     pairs.push_back({key, folly::to<std::string>(bucket)});
   };
   auto ctx = ProxyRequestContextWithInfo<RouterInfo>::createRecording(