move source ip from req context to req common

Summary: - move source ip to req common - use folly::IPAddress instead of std::string - change callsites to get ip from req rather than ctx Reviewed By: disylh Differential Revision: D64721214 fbshipit-source-id: d5d7bc35def6555e5de78db50c704b9bed0e77ee
facebook · Nov 8, 2024 · 4450ce9 · 4450ce9
1 parent ebb0407
commit 4450ce9
Show file tree

Hide file tree

Showing 14 changed files with 138 additions and 75 deletions.
diff --git a/mcrouter/CarbonRouterClient-inl.h b/mcrouter/CarbonRouterClient-inl.h
@@ -94,13 +94,9 @@ bool srHostInfoPtrFuncCarbonRouterClient(
 
 template <class RouterInfo>
 template <class Request, class F>
-bool CarbonRouterClient<RouterInfo>::send(
-    const Request& req,
-    F&& callback,
-    folly::StringPiece ipAddr) {
-  auto makePreq = [this, ipAddr, &req, &callback](bool inBatch) mutable {
-    return makeProxyRequestContext(
-        req, std::forward<F>(callback), ipAddr, inBatch);
+bool CarbonRouterClient<RouterInfo>::send(const Request& req, F&& callback) {
+  auto makePreq = [this, &req, &callback](bool inBatch) mutable {
+    return makeProxyRequestContext(req, std::forward<F>(callback), inBatch);
   };
 
   auto cancelRemaining = [&req, &callback]() {
@@ -259,15 +255,14 @@ template <class InputIt, class F>
 bool CarbonRouterClient<RouterInfo>::send(
     InputIt begin,
     InputIt end,
-    F&& callback,
-    folly::StringPiece ipAddr) {
+    F&& callback) {
   using IterReference = typename std::iterator_traits<InputIt>::reference;
   using Request = typename std::decay<decltype(detail::unwrapRequest(
       std::declval<IterReference>()))>::type;
 
-  auto makeNextPreq = [this, ipAddr, &callback, &begin](bool inBatch) {
+  auto makeNextPreq = [this, &callback, &begin](bool inBatch) {
     auto proxyRequestContext = makeProxyRequestContext(
-        detail::unwrapRequest(*begin), callback, ipAddr, inBatch);
+        detail::unwrapRequest(*begin), callback, inBatch);
     ++begin;
     return proxyRequestContext;
   };
@@ -371,7 +366,6 @@ std::unique_ptr<ProxyRequestContextWithInfo<RouterInfo>>
 CarbonRouterClient<RouterInfo>::makeProxyRequestContext(
     const Request& req,
     CallbackFunc&& callback,
-    folly::StringPiece ipAddr,
     bool inBatch) {
   Proxy<RouterInfo>* proxy = proxies_[proxyIdx_];
   if (mode_ == ThreadMode::AffinitizedRemoteThread) {
@@ -398,9 +392,6 @@ CarbonRouterClient<RouterInfo>::makeProxyRequestContext(
       });
 
   proxyRequestContext->setRequester(self_);
-  if (!ipAddr.empty()) {
-    proxyRequestContext->setUserIpAddress(ipAddr);
-  }
   return proxyRequestContext;
 }
 

diff --git a/mcrouter/CarbonRouterClient.h b/mcrouter/CarbonRouterClient.h
@@ -100,10 +100,7 @@ class CarbonRouterClient : public CarbonRouterClientBase {
    *       callback is called.
    */
   template <class Request, class F>
-  bool send(
-      const Request& req,
-      F&& callback,
-      folly::StringPiece ipAddr = folly::StringPiece());
+  bool send(const Request& req, F&& callback);
 
   /**
    * Multi requests version of send.
@@ -122,11 +119,7 @@ class CarbonRouterClient : public CarbonRouterClientBase {
    *       each request.
    */
   template <class InputIt, class F>
-  bool send(
-      InputIt begin,
-      InputIt end,
-      F&& callback,
-      folly::StringPiece ipAddr = folly::StringPiece());
+  bool send(InputIt begin, InputIt end, F&& callback);
 
   CacheClientCounters getStatCounters() noexcept {
     return stats_.getCounters();
@@ -222,7 +215,6 @@ class CarbonRouterClient : public CarbonRouterClientBase {
    * @param req         The request to be routed.
    * @param callback    The callback function to be called once the reply
    *                    is received.
-   * @param ipAddr      The ip address of the caller (can be empty).
    * @param inBatch     Whether or not the given request is part of a batch
    *                    of requests.
    *
@@ -233,7 +225,6 @@ class CarbonRouterClient : public CarbonRouterClientBase {
   makeProxyRequestContext(
       const Request& req,
       CallbackFunc&& callback,
-      folly::StringPiece ipAddr,
       bool inBatch);
 
   bool shouldDelayNotification(size_t batchSize) const;

diff --git a/mcrouter/ProxyRequestContext.h b/mcrouter/ProxyRequestContext.h
@@ -123,14 +123,6 @@ class ProxyRequestContext {
         "of ProxyRequestContext");
   }
 
-  const std::string& userIpAddress() const noexcept {
-    return userIpAddr_;
-  }
-
-  void setUserIpAddress(folly::StringPiece newAddr) noexcept {
-    userIpAddr_ = newAddr.str();
-  }
-
   bool isProcessing() const {
     return processing_;
   }
@@ -219,8 +211,6 @@ class ProxyRequestContext {
 
   uint64_t senderIdForTest_{0};
 
-  std::string userIpAddr_;
-
   ProxyRequestPriority priority_{ProxyRequestPriority::kCritical};
 
   bool failoverDisabled_{false};

diff --git a/mcrouter/ServerOnRequest.h b/mcrouter/ServerOnRequest.h
@@ -200,12 +200,11 @@ class ServerOnRequest {
       }
     };
 
-    folly::Optional<std::string> peerIp;
-    if (retainSourceIp_ && (peerIp = ctxRef.getPeerSocketAddressStr())) {
-      client_.send(reqRef, std::move(cb), *peerIp);
-    } else {
-      client_.send(reqRef, std::move(cb));
+    std::optional<folly::IPAddress> peerIp;
+    if (retainSourceIp_ && (peerIp = ctxRef.getPeerSocketIPAddress())) {
+      reqRef.setSourceIpAddr(*peerIp);
     }
+    client_.send(reqRef, std::move(cb));
   }
 
  private:

diff --git a/mcrouter/lib/HashSelector.h b/mcrouter/lib/HashSelector.h
@@ -14,6 +14,7 @@
 
 #include "mcrouter/McrouterFiberContext.h"
 #include "mcrouter/lib/HashUtil.h"
+#include "mcrouter/lib/fbi/cpp/util.h"
 
 namespace facebook {
 namespace memcache {
@@ -88,17 +89,16 @@ class BucketHashSelector : public HashSelectorBase<HashFunc> {
         clientFanout_(clientFanout) {}
 
   template <class Request>
-  size_t select(const Request& /*req*/, size_t size) const {
+  size_t select(const Request& req, size_t size) const {
     auto bucketId = mcrouter::fiber_local<RouterInfo>::getBucketId();
     checkRuntime(
         bucketId.has_value(),
         "The context doesn't contain bucket id. You must use McBucketRoute in front of bucketized PoolRoute");
     // Hash functions can be stack-intensive, so jump back to the main context
     auto key = folly::to<std::string>(*bucketId);
     if (this->clientFanout_) {
-      auto& ctx = mcrouter::fiber_local<RouterInfo>::getSharedCtx();
-      if (ctx) {
-        key += ctx->userIpAddress();
+      if (auto sourceIpAddr = req.getSourceIpAddr()) {
+        key += sourceIpAddr->str();
       }
     }
     return folly::fibers::runInMainContext(

diff --git a/mcrouter/lib/carbon/RequestCommon.h b/mcrouter/lib/carbon/RequestCommon.h
@@ -10,6 +10,7 @@
 #include <optional>
 #include <utility>
 
+#include <folly/IPAddress.h>
 #include <folly/io/IOBuf.h>
 
 #include "mcrouter/lib/carbon/MessageCommon.h"
@@ -93,6 +94,14 @@ class RequestCommon : public MessageCommon {
     clientIdentifier_ = clientIdentifier.str();
   }
 
+  const std::optional<folly::IPAddress>& getSourceIpAddr() const noexcept {
+    return sourceIpAddr_;
+  }
+
+  void setSourceIpAddr(const folly::IPAddress& sourceIpAddr) noexcept {
+    sourceIpAddr_ = sourceIpAddr;
+  }
+
  protected:
   void markBufferAsDirty() {
     serializedBuffer_ = nullptr;
@@ -104,6 +113,8 @@ class RequestCommon : public MessageCommon {
   std::optional<std::string> cryptoAuthToken_;
   // Hash string of primary (non-host) tls client identities
   std::optional<std::string> clientIdentifier_;
+  // Source ip address.
+  std::optional<folly::IPAddress> sourceIpAddr_;
 };
 
 } // namespace carbon
diff --git a/mcrouter/lib/network/McServerRequestContext.cpp b/mcrouter/lib/network/McServerRequestContext.cpp
@@ -91,12 +91,13 @@ ServerLoad McServerRequestContext::getServerLoad() const noexcept {
   return ServerLoad::zero();
 }
 
-folly::Optional<std::string> McServerRequestContext::getPeerSocketAddressStr() {
-  folly::Optional<std::string> peerAddressStr;
+std::optional<folly::IPAddress>
+McServerRequestContext::getPeerSocketIPAddress() {
+  std::optional<folly::IPAddress> peerAddress;
   if (session_) {
-    peerAddressStr = session_->getSocketAddress().getAddressStr();
+    peerAddress = session_->getSocketAddress().getIPAddress();
   }
-  return peerAddressStr;
+  return peerAddress;
 }
 
 folly::Optional<struct sockaddr_storage>

diff --git a/mcrouter/lib/network/McServerRequestContext.h b/mcrouter/lib/network/McServerRequestContext.h
@@ -64,7 +64,7 @@ class McServerRequestContext {
   ServerLoad getServerLoad() const noexcept;
 
   folly::Optional<struct sockaddr_storage> getPeerSocketAddress();
-  folly::Optional<std::string> getPeerSocketAddressStr();
+  std::optional<folly::IPAddress> getPeerSocketIPAddress();
 
   folly::EventBase& getSessionEventBase() const noexcept;
   const apache::thrift::Cpp2RequestContext* getThriftRequestContext()

diff --git a/mcrouter/lib/network/McThriftContext.h b/mcrouter/lib/network/McThriftContext.h
@@ -32,13 +32,13 @@ class McThriftContext {
     ctx.underlying_->exception(std::move(ex));
   }
 
-  folly::Optional<std::string> getPeerSocketAddressStr() {
-    folly::Optional<std::string> peerAddressStr;
+  std::optional<folly::IPAddress> getPeerSocketIPAddress() {
+    std::optional<folly::IPAddress> peerIPAddr;
     auto connectionCxt = underlying_->getConnectionContext();
     if (connectionCxt) {
-      peerAddressStr = connectionCxt->getPeerAddress()->getAddressStr();
+      peerIPAddr = connectionCxt->getPeerAddress()->getIPAddress();
     }
-    return peerAddressStr;
+    return peerIPAddr;
   }
 
   folly::Optional<struct sockaddr_storage> getPeerSocketAddress() {

diff --git a/mcrouter/lib/test/RouteHandleTest.cpp b/mcrouter/lib/test/RouteHandleTest.cpp
@@ -347,6 +347,47 @@ TEST(routeHandleTest, hashNoSalt) {
   });
 }
 
+TEST(routeHandleTest, bucketHashSelector) {
+  vector<std::shared_ptr<TestHandle>> test_handles{
+      make_shared<TestHandle>(GetRouteTestData(carbon::Result::FOUND, "a")),
+      make_shared<TestHandle>(GetRouteTestData(carbon::Result::FOUND, "b")),
+      make_shared<TestHandle>(GetRouteTestData(carbon::Result::FOUND, "c")),
+  };
+  auto outOfRangeRh = createNullRoute<typename TestRouterInfo::RouteHandleIf>();
+
+  TestFiberManager<TestRouterInfo> fm;
+
+  TestRouteHandle<SelectionRoute<
+      TestRouterInfo,
+      BucketHashSelector<HashFunc, TestRouterInfo>>>
+      rh(get_route_handles(test_handles),
+         BucketHashSelector<HashFunc, TestRouterInfo>(
+             /* salt= */ "1",
+             HashFunc(test_handles.size()),
+             /*clientFanout=*/true),
+         std::move(outOfRangeRh));
+  McGetRequest getReq("0");
+  fm.run([&]() {
+    fiber_local<TestRouterInfo>::setBucketId(1);
+    auto reply = rh.route(getReq);
+    EXPECT_EQ("c", carbon::valueRangeSlow(reply).str());
+  });
+
+  fm.run([&]() {
+    fiber_local<TestRouterInfo>::setBucketId(1);
+    getReq.setSourceIpAddr(folly::IPAddress("0.0.0.0"));
+    auto reply = rh.route(getReq);
+    EXPECT_EQ("b", carbon::valueRangeSlow(reply).str());
+  });
+
+  fm.run([&]() {
+    fiber_local<TestRouterInfo>::setBucketId(1);
+    getReq.setSourceIpAddr(folly::IPAddress("1.0.0.0"));
+    auto reply = rh.route(getReq);
+    EXPECT_EQ("c", carbon::valueRangeSlow(reply).str());
+  });
+}
+
 TEST(routeHandleTest, hashSalt) {
   vector<std::shared_ptr<TestHandle>> test_handles{
       make_shared<TestHandle>(GetRouteTestData(carbon::Result::FOUND, "a")),

diff --git a/mcrouter/routes/LoggingRoute.h b/mcrouter/routes/LoggingRoute.h
@@ -67,12 +67,11 @@ class LoggingRoute {
 
     // Pull the IP (if available) out of the saved request
     auto& ctx = mcrouter::fiber_local<RouterInfo>::getSharedCtx();
-    auto& ip = ctx->userIpAddress();
-    folly::StringPiece userIp;
-    if (!ip.empty()) {
-      userIp = ip;
+    std::string userIpStr;
+    if (auto sourceIpAddr = req.getSourceIpAddr()) {
+      userIpStr = sourceIpAddr->str();
     } else {
-      userIp = "N/A";
+      userIpStr = "N/A";
     }
 
     auto& callback = ctx->proxy().router().postprocessCallback();
@@ -83,7 +82,7 @@ class LoggingRoute {
             carbon::convertToFollyDynamic(req),
             carbon::convertToFollyDynamic(reply),
             Request::name,
-            userIp);
+            userIpStr);
       }
     }
     return reply;

diff --git a/mcrouter/routes/OriginalClientHashRoute.h b/mcrouter/routes/OriginalClientHashRoute.h
@@ -55,11 +55,10 @@ class OriginalClientHashRoute {
       return false;
     }
 
-    auto& ctx = mcrouter::fiber_local<RouterInfo>::getSharedCtx();
-    auto& ip = ctx->userIpAddress();
-    if (!ip.empty()) {
+    if (auto ip = req.getSourceIpAddr()) {
       return t(
-          *children_[(folly::Hash()(ip) + offset_) % children_.size()], req);
+          *children_[(folly::Hash()(ip->str()) + offset_) % children_.size()],
+          req);
     }
     return false;
   }
@@ -69,11 +68,9 @@ class OriginalClientHashRoute {
     ReplyT<Request> reply;
     assert(children_.size() > 1);
 
-    auto& ctx = mcrouter::fiber_local<RouterInfo>::getSharedCtx();
-    auto& ip = ctx->userIpAddress();
-    if (!ip.empty()) {
-      return children_[(folly::Hash()(ip) + offset_) % children_.size()]->route(
-          req);
+    if (auto ip = req.getSourceIpAddr()) {
+      return children_[(folly::Hash()(ip->str()) + offset_) % children_.size()]
+          ->route(req);
     } else {
       return createReply<Request>(
           ErrorReply, carbon::Result::LOCAL_ERROR, "Missing IP");

diff --git a/mcrouter/routes/test/LoggingRouteTest.cpp b/mcrouter/routes/test/LoggingRouteTest.cpp
@@ -0,0 +1,42 @@
+/*
+ * 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 <memory>
+#include <string>
+
+#include <gtest/gtest.h>
+
+#include "mcrouter/lib/test/TestRouteHandle.h"
+#include "mcrouter/routes/LoggingRoute.h"
+#include "mcrouter/routes/test/RouteHandleTestUtil.h"
+
+using TestRouterInfo = facebook::memcache::TestRouterInfo;
+using TestHandle =
+    facebook::memcache::TestHandleImpl<facebook::memcache::TestRouteHandleIf>;
+using LoggingRoute = facebook::memcache::mcrouter::LoggingRoute<TestRouterInfo>;
+
+TEST(LoggingRouteTest, basic) {
+  std::shared_ptr<TestHandle> test_handle = std::make_shared<TestHandle>(
+      facebook::memcache::GetRouteTestData(carbon::Result::FOUND, "a"));
+
+  auto rh = facebook::memcache::mcrouter::createLoggingRoute<TestRouterInfo>(
+      test_handle->rh);
+  facebook::memcache::McGetRequest req("key");
+  req.setSourceIpAddr(folly::IPAddress("1.2.3.4"));
+
+  facebook::memcache::TestFiberManager<facebook::memcache::TestRouterInfo> fm;
+  fm.run([&]() {
+    facebook::memcache::mcrouter::mockFiberContext();
+    facebook::memcache::mcrouter::getTestRouter()->setPostprocessCallback(
+        [](const folly::dynamic&,
+           const folly::dynamic&,
+           const char* const,
+           const folly::StringPiece ip) { EXPECT_EQ("1.2.3.4", ip); });
+    auto reply = rh->route(req);
+    EXPECT_EQ("a", carbon::valueRangeSlow(reply).str());
+  });
+}