Add option to wrap srroute with refill route

Summary: as title

Reviewed By: disylh, stuclar

Differential Revision: D61394564

fbshipit-source-id: 3c9d95a47667a532b4a262b32c119ba85e2adb2c

mcrouter/routes/McRefillRoute.h

@@ -170,6 +170,14 @@ typename RouterInfo::RouteHandlePtr makeMcRefillRoute(
       factory.create(json["primary"]), factory.create(json["refill"]));
 }
 
+template <class RouterInfo>
+typename RouterInfo::RouteHandlePtr makeMcRefillRouteUsePrimaryAndRefill(
+    typename RouterInfo::RouteHandlePtr primary,
+    typename RouterInfo::RouteHandlePtr refill) {
+  return makeRouteHandleWithInfo<RouterInfo, McRefillRoute>(
+      std::move(primary), std::move(refill));
+}
+
 } // namespace mcrouter
 } // namespace memcache
 } // namespace facebook

mcrouter/routes/McRouteHandleProvider-inl.h

@@ -34,6 +34,7 @@
 #include "mcrouter/routes/FailoverRoute.h"
 #include "mcrouter/routes/HashRouteFactory.h"
 #include "mcrouter/routes/McBucketRoute.h"
+#include "mcrouter/routes/McRefillRoute.h"
 #include "mcrouter/routes/PoolRouteUtils.h"
 #include "mcrouter/routes/RateLimitRoute.h"
 #include "mcrouter/routes/RateLimiter.h"
@@ -581,6 +582,26 @@ McRouteHandleProvider<RouterInfo>::createSRRoute(
     route = std::move(makeShadowRoutes(
         factory, json, {std::move(route)}, proxy_, *extraProvider_)[0]);
   }
+
+  if constexpr (RouterInfo::hasMcRefillRoute) {
+    bool refillOnMiss = false;
+    if (auto* jRefillOnMiss = json.get_ptr("refill_on_miss")) {
+      refillOnMiss = parseBool(*jRefillOnMiss, "refill_on_miss");
+    }
+    if (refillOnMiss) {
+      auto jRefillFromTier = json.get_ptr("refill_from_tier");
+      checkLogic(
+          jRefillFromTier != nullptr,
+          "SRroute: 'refill_from_tier' property is missing");
+      folly::dynamic refillJson = json;
+      refillJson["service_name"] = *jRefillFromTier;
+      refillJson["type"] = "SRRoute";
+
+      route = makeMcRefillRouteUsePrimaryAndRefill<RouterInfo>(
+          route, factoryFunc(factory, refillJson, proxy_));
+    }
+  }
+
   route = bucketize(std::move(route), json);
 
   if (needAsynclog) {

mcrouter/test/cpp_unit_tests/mc_route_handle_provider_test.cpp

@@ -81,6 +81,20 @@ const char* const kBucketizedPoolRoute =
    "bucketization_keyspace": "tst"
 })";
 
+const char* const kRefillOnMissSRRoute =
+    R"({
+   "type": "SRRoute",
+   "service_name": "twmemcache.shadow.bucketization-test",
+   "server_timeout": 200,
+   "asynclog_name": "test.asynclog",
+   "axonlog": false,
+   "bucketize": true,
+   "total_buckets": 1000,
+   "bucketization_keyspace": "tst",
+   "refill_on_miss": true,
+   "refill_from_tier": "twmemcache.shadow.mcrouter-rc"
+})";
+
 struct TestSetup {
  public:
   TestSetup()
@@ -176,3 +190,12 @@ TEST(McRouteHandleProvider, bucketized_pool_route_and_mcreplay_asynclogRoutes) {
       "bucketize|total_buckets=1000|bucketization_keyspace=tst|prefix_map_enabled=false",
       asynclogRoutes["test.asynclog"]->routeName());
 }
+
+TEST(McRouteHandleProvider, bucketized_sr_route_with_refill_on_miss) {
+  TestSetup setup;
+  auto rh = setup.getRoute(kRefillOnMissSRRoute);
+  EXPECT_TRUE(rh != nullptr);
+  EXPECT_EQ(
+      "bucketize|total_buckets=1000|bucketization_keyspace=tst|prefix_map_enabled=false",
+      rh->routeName());
+}