diff --git a/mcrouter/ProxyRequestContext.h b/mcrouter/ProxyRequestContext.h index 52b37ea8c..dbe8a7c20 100644 --- a/mcrouter/ProxyRequestContext.h +++ b/mcrouter/ProxyRequestContext.h @@ -50,8 +50,11 @@ class ProxyRequestContext { using ClientCallback = std::function; using ShardSplitCallback = std::function; - using BucketizationCallback = - std::function; + using BucketizationCallback = std::function)>; virtual ~ProxyRequestContext(); @@ -80,13 +83,15 @@ class ProxyRequestContext { } } + template 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 index 1da23d484..48e04cb93 100644 --- a/mcrouter/ServiceInfo-inl.h +++ b/mcrouter/ServiceInfo-inl.h @@ -221,8 +221,8 @@ class GetBucketCommandDispatcher { folly::StringPiece keyStr, Proxy& proxy, const ProxyRoute& proxyRoute) { - using KeyBucketKeyspaceTuple = - std::tuple; + using KeyBucketKeyspaceTuple = std:: + tuple>; proxy.fiberManager().addTaskFinally( [keyStr, &proxy, &proxyRoute]() { auto keyBucketKeyspaceTuples = @@ -230,11 +230,13 @@ class GetBucketCommandDispatcher { auto cb = [&keyBucketKeyspaceTuples]( std::string keyRecorded, const uint64_t bucketIdRecorded, - const std::string_view bucketizationKeyspaceRecorded) { + const std::string_view bucketizationKeyspaceRecorded, + std::optional usecaseId) { keyBucketKeyspaceTuples->push_back(std::make_tuple( std::move(keyRecorded), folly::to(bucketIdRecorded), - folly::to(bucketizationKeyspaceRecorded))); + folly::to(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(*usecaseId) + : "null"); } ReplyT 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 index 57c6ecef0..158b98908 100644 --- a/mcrouter/lib/network/MessageHelpers.h +++ b/mcrouter/lib/network/MessageHelpers.h @@ -217,8 +217,7 @@ class HasUsecaseIdTrait : public std::false_type {}; template class HasUsecaseIdTrait< Message, - std::void_t< - decltype(std::declval>().usecaseId_ref())>> + std::void_t().usecaseId_ref())>> : public std::true_type {}; template diff --git a/mcrouter/routes/McBucketRoute.h b/mcrouter/routes/McBucketRoute.h index 860bd1a4f..d796590c4 100644 --- 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::runWithLocals([this, &req, &t, bucketId]() { fiber_local::setBucketId(bucketId); @@ -120,7 +121,8 @@ class McBucketRoute { ctx->recordBucketizationData( req.key_ref()->keyWithoutRoute().str(), bucketId, - bucketizationKeyspace_); + bucketizationKeyspace_, + req); return createReply(DefaultReply, req); } return routeImpl(req, bucketId); diff --git a/mcrouter/routes/test/RouteHandleTestUtil.h b/mcrouter/routes/test/RouteHandleTestUtil.h index bce831272..60d1faa6e 100644 --- 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 /* usecaseId */) { pairs.push_back({key, folly::to(bucket)}); }; auto ctx = ProxyRequestContextWithInfo::createRecording(