Support route handles that must see every request

Summary:
A binary that embeds mcrouter sometimes needs a route handle that sees *every*
request — an interceptor, a local overlay, a shim in front of production. Today
the only way to get one is to author or rewrite the routing config, which means
knowing how that config expresses its entry point (`route` vs `routes`, alias
lists, `PrefixSelectorRoute` policies), reproducing whatever the config would
otherwise have routed to, and re-doing all of it whenever that config changes.
Configs that reach their entry point through JSONM macros cannot be edited that
way at all, since the leaves do not exist until the preprocessor has run.

Give binaries a way to say "put this on top" and let mcrouter resolve the config
normally.

`ExtraRouteHandleProviderIf::wrapRoot()` is a new hook, defaulting to identity,
called from `ProxyRoute`'s constructor once the root handle is built. It is
applied after the existing `BigValueRoute` and `LoggingRoute` wrappers and above
routing prefix selection, so the injected handle sees every request and is
created exactly once per proxy, whatever shape the loaded config has.
`McRouteHandleProvider::extraProvider()` exposes the provider so `ProxyConfig`
can pass it to `ProxyRoute`, alongside the `release*()` accessors it already
calls on that object a few lines earlier.

Nothing changes for existing routers: the `ProxyRoute` parameter defaults to
`nullptr`, and no in-tree provider overrides `wrapRoot()` in this diff, so every
current mcrouter builds an identical route handle tree.

The first user is TAO Sandbox's `mock_tao` / `mock_ucache` (next in this stack).
It needs `MockTaoRoute` / `MockUcacheRoute` in front of whatever config the
flavor and the host's `/etc/mcrouter` overrides resolve to; today it hand-writes
a config that bypasses both and has drifted from the real one as a result.

Reviewed By: ghostonhuang

Differential Revision: D116661224

fbshipit-source-id: e07f7a2f69ca3e95cc35aaf4275d126841ab8964

Author: Vrishal Kulkarni

mcrouter/ProxyConfig-inl.h

@@ -152,7 +152,8 @@ ProxyConfig<RouterInfo>::ProxyConfig(
           .enableAsyncDlBroadcast = enableAsyncDlBroadcast,
           .enableSetDistribution = enableSetDistribution,
           .enableCrossRegionSetRpc = enableCrossRegionSetRpc,
-          .enableGlobalBigValueRoute = enableGlobalBigValueRoute});
+          .enableGlobalBigValueRoute = enableGlobalBigValueRoute},
+      provider.extraProvider());
   serviceInfo_ = std::make_shared<ServiceInfo<RouterInfo>>(proxy, *this);
 }
 

mcrouter/routes/ExtraRouteHandleProviderIf.h

@@ -43,6 +43,20 @@ class ExtraRouteHandleProviderIf {
       folly::StringPiece type,
       const folly::dynamic& json) = 0;
 
+  /**
+   * Wraps the root of the route handle tree, after the config's entry point
+   * has been turned into route handles.
+   *
+   * Lets a binary inject a route handle that must see every request no matter
+   * how the config expresses its entry point (route vs routes, aliases,
+   * prefix selectors). The wrapper sits above routing prefix selection, so it
+   * is created exactly once per proxy.
+   */
+  virtual std::shared_ptr<RouteHandleIf> wrapRoot(
+      std::shared_ptr<RouteHandleIf> root) {
+    return root;
+  }
+
   virtual ~ExtraRouteHandleProviderIf() {}
 };
 } // namespace mcrouter

mcrouter/routes/McRouteHandleProvider.h

@@ -146,6 +146,10 @@ class McRouteHandleProvider
     return std::move(accessPoints_);
   }
 
+  ExtraRouteHandleProviderIf<RouterInfo>* extraProvider() {
+    return extraProvider_.get();
+  }
+
   ~McRouteHandleProvider() override;
 
  private:

mcrouter/routes/ProxyRoute-inl.h

@@ -12,6 +12,7 @@
 #include "mcrouter/Proxy.h"
 #include "mcrouter/lib/config/RouteHandleBuilder.h"
 #include "mcrouter/routes/BigValueRoute.h"
+#include "mcrouter/routes/ExtraRouteHandleProviderIf.h"
 #include "mcrouter/routes/LoggingRoute.h"
 #include "mcrouter/routes/McRouteHandleBuilder.h"
 
@@ -42,7 +43,8 @@ template <class RouterInfo>
 ProxyRoute<RouterInfo>::ProxyRoute(
     Proxy<RouterInfo>& proxy,
     const RouteSelectorMap<typename RouterInfo::RouteHandleIf>& routeSelectors,
-    RootRouteRolloutOpts rolloutOpts)
+    RootRouteRolloutOpts rolloutOpts,
+    ExtraRouteHandleProviderIf<RouterInfo>* extraProvider)
     : proxy_(proxy),
       root_(
           makeRouteHandleWithInfo<RouterInfo, RootRoute>(
@@ -57,6 +59,11 @@ ProxyRoute<RouterInfo>::ProxyRoute(
   if (proxy_.getRouterOptions().enable_logging_route) {
     root_ = createLoggingRoute<RouterInfo>(std::move(root_));
   }
+  // Applied last so the injected handle sees every request, before routing
+  // prefix selection and before any of the wrappers above.
+  if (extraProvider) {
+    root_ = extraProvider->wrapRoot(std::move(root_));
+  }
 }
 
 template <class RouterInfo>

mcrouter/routes/ProxyRoute.h

@@ -31,6 +31,9 @@ namespace mcrouter {
 template <class RouterInfo>
 class Proxy;
 
+template <class RouterInfo>
+class ExtraRouteHandleProviderIf;
+
 /**
  * This is the top-most level of Mcrouter's RouteHandle tree.
  */
@@ -45,7 +48,8 @@ class ProxyRoute {
       Proxy<RouterInfo>& proxy,
       const RouteSelectorMap<typename RouterInfo::RouteHandleIf>&
           routeSelectors,
-      RootRouteRolloutOpts rolloutOpts);
+      RootRouteRolloutOpts rolloutOpts,
+      ExtraRouteHandleProviderIf<RouterInfo>* extraProvider = nullptr);
 
   template <class Request>
   bool traverse(

mcrouter/test/cpp_unit_tests/mc_route_handle_provider_test.cpp

@@ -15,11 +15,17 @@
 #include "mcrouter/CarbonRouterInstance.h"
 #include "mcrouter/PoolFactory.h"
 #include "mcrouter/Proxy.h"
+#include "mcrouter/lib/RouteHandleTraverser.h"
+#include "mcrouter/lib/config/RouteHandleBuilder.h"
 #include "mcrouter/lib/config/RouteHandleFactory.h"
 #include "mcrouter/lib/network/gen/MemcacheRouterInfo.h"
 #include "mcrouter/options.h"
+#include "mcrouter/routes/McExtraRouteHandleProvider.h"
 #include "mcrouter/routes/McRouteHandleProvider.h"
 #include "mcrouter/routes/McrouterRouteHandle.h"
+#include "mcrouter/routes/PrefixSelectorRoute.h"
+#include "mcrouter/routes/ProxyRoute.h"
+#include "mcrouter/routes/RouteSelectorMap.h"
 
 using namespace facebook::memcache;
 using namespace facebook::memcache::mcrouter;
@@ -113,6 +119,14 @@ struct TestSetup {
     return rhProvider_;
   }
 
+  RouteHandleFactory<McrouterRouteHandleIf>& factory() {
+    return rhFactory_;
+  }
+
+  Proxy<MemcacheRouterInfo>& proxy() {
+    return *router_->getProxy(0);
+  }
+
   McrouterRouteHandlePtr getRoute(const char* jsonStr) {
     return rhFactory_.create(parseJsonString(jsonStr));
   }
@@ -131,6 +145,63 @@ struct TestSetup {
   }
 };
 
+// Recognizable pass-through handle, so a traversal can tell whether the root
+// of the tree is the one ProxyRoute built or one an extra provider put there.
+template <class RouteHandleIf>
+class MarkerRoute {
+ public:
+  static std::string routeName() {
+    return "marker";
+  }
+
+  explicit MarkerRoute(std::shared_ptr<RouteHandleIf> child)
+      : child_(std::move(child)) {}
+
+  template <class Request>
+  bool traverse(const Request& req, RouteHandleTraverser<RouteHandleIf>& t)
+      const {
+    return t(*child_, req);
+  }
+
+  template <class Request>
+  ReplyT<Request> route(const Request& req) const {
+    return child_->route(req);
+  }
+
+ private:
+  std::shared_ptr<RouteHandleIf> child_;
+};
+
+class MarkerExtraProvider
+    : public McExtraRouteHandleProvider<MemcacheRouterInfo> {
+ public:
+  McrouterRouteHandlePtr wrapRoot(McrouterRouteHandlePtr root) override {
+    return makeRouteHandle<McrouterRouteHandleIf, MarkerRoute>(std::move(root));
+  }
+};
+
+// Single-entry selector map, enough to build a ProxyRoute.
+RouteSelectorMap<McrouterRouteHandleIf> nullRouteSelectors(TestSetup& setup) {
+  RouteSelectorMap<McrouterRouteHandleIf> selectors;
+  selectors[setup.proxy().getRouterOptions().default_route] =
+      std::make_shared<PrefixSelectorRoute<McrouterRouteHandleIf>>(
+          setup.factory(), parseJsonString(R"("NullRoute")"));
+  return selectors;
+}
+
+// Name of the first route handle a traversal of `proxyRoute` reaches.
+std::string rootRouteName(ProxyRoute<MemcacheRouterInfo>& proxyRoute) {
+  std::string name;
+  RouteHandleTraverser<McrouterRouteHandleIf> t{
+      [&name](const McrouterRouteHandleIf& rh) {
+        if (name.empty()) {
+          name = rh.routeName();
+        }
+      }};
+  proxyRoute.traverse(McGetRequest("key"), t);
+  return name;
+}
+
 } // namespace
 
 TEST(McRouteHandleProviderTest, sanity) {
@@ -210,3 +281,31 @@ TEST(McRouteHandleProvider, bucketized_pool_route_and_mcreplay_asynclogRoutes) {
       "bucketize|total_buckets=1000|bucketization_keyspace=tst|prefix_map_enabled=false",
       asynclogRoutes["test.asynclog"]->routeName());
 }
+
+// A provider that does not override wrapRoot leaves the tree alone, so
+// existing routers keep the root that ProxyRoute built for them.
+TEST(McRouteHandleProvider, wrap_root_defaults_to_identity) {
+  TestSetup setup;
+  auto rh = setup.getRoute(kConstShard);
+  ASSERT_TRUE(setup.provider().extraProvider() != nullptr);
+  EXPECT_EQ(rh, setup.provider().extraProvider()->wrapRoot(rh));
+}
+
+TEST(McRouteHandleProvider, proxy_route_without_provider_is_unwrapped) {
+  TestSetup setup;
+  auto selectors = nullRouteSelectors(setup);
+  ProxyRoute<MemcacheRouterInfo> proxyRoute(
+      setup.proxy(), selectors, RootRouteRolloutOpts{});
+  EXPECT_EQ("root", rootRouteName(proxyRoute));
+}
+
+// The handle an extra provider returns from wrapRoot ends up above the root,
+// so it sees every request no matter how the config selects routes.
+TEST(McRouteHandleProvider, proxy_route_applies_wrap_root) {
+  TestSetup setup;
+  auto selectors = nullRouteSelectors(setup);
+  MarkerExtraProvider extraProvider;
+  ProxyRoute<MemcacheRouterInfo> proxyRoute(
+      setup.proxy(), selectors, RootRouteRolloutOpts{}, &extraProvider);
+  EXPECT_EQ("marker", rootRouteName(proxyRoute));
+}