RoutePolicyMap V2 (#433)

Summary:
Pull Request resolved: #433

RoutePolicyMap implemented on top of LowerBoundPrefixmap

Reviewed By: disylh

Differential Revision: D49188195

fbshipit-source-id: 32b2769531c41ba65739719d2a7925ec260cb40d

mcrouter/lib/fbi/cpp/LowerBoundPrefixMap.cpp

@@ -25,7 +25,7 @@ std::ostream& operator<<(std::ostream& os, const SmallPrefix& self) {
 }
 
 LowerBoundPrefixMapCommon::LowerBoundPrefixMapCommon(
-    std::span<const std::string_view> sortedUniquePrefixes) {
+    const std::vector<std::string_view>& sortedUniquePrefixes) {
   smallPrefixes_.reserve(sortedUniquePrefixes.size() + 1);
   markers_.reserve(sortedUniquePrefixes.size() + 1);
   previousPrefix_.reserve(sortedUniquePrefixes.size());

mcrouter/lib/fbi/cpp/LowerBoundPrefixMap.h

@@ -8,16 +8,15 @@
 #pragma once
 
 #include <folly/CPortability.h>
+#include <folly/Range.h>
 #include <folly/container/Iterator.h>
 #include <folly/lang/Bits.h>
 #include <folly/sorted_vector_types.h>
 
 #include <algorithm>
-#include <compare>
 #include <iosfwd>
 #include <limits>
 #include <optional>
-#include <span>
 #include <string>
 #include <string_view>
 #include <utility>
@@ -45,8 +44,38 @@ class SmallPrefix {
     data_ = folly::Endian::swap(data_);
   }
 
-  bool operator==(const SmallPrefix&) const = default;
-  auto operator<=>(const SmallPrefix&) const = default;
+  // default operator== and operator<=> are ok here
+  // but not everything is C++20.
+
+  FOLLY_ALWAYS_INLINE
+  friend bool operator==(const SmallPrefix& x, const SmallPrefix& y) {
+    return x.data_ == y.data_;
+  }
+
+  FOLLY_ALWAYS_INLINE
+  friend bool operator!=(const SmallPrefix& x, const SmallPrefix& y) {
+    return !(x == y);
+  }
+
+  FOLLY_ALWAYS_INLINE
+  friend bool operator<(const SmallPrefix& x, const SmallPrefix& y) {
+    return x.data_ < y.data_;
+  }
+
+  FOLLY_ALWAYS_INLINE
+  friend bool operator<=(const SmallPrefix& x, const SmallPrefix& y) {
+    return !(y < x);
+  }
+
+  FOLLY_ALWAYS_INLINE
+  friend bool operator>=(const SmallPrefix& x, const SmallPrefix& y) {
+    return !(x < y);
+  }
+
+  FOLLY_ALWAYS_INLINE
+  friend bool operator>(const SmallPrefix& x, const SmallPrefix& y) {
+    return y < x;
+  }
 
   friend std::ostream& operator<<(std::ostream& os, const SmallPrefix& self);
 
@@ -72,7 +101,7 @@ class SmallPrefix {
 struct LowerBoundPrefixMapCommon {
   LowerBoundPrefixMapCommon() = default;
   explicit LowerBoundPrefixMapCommon(
-      std::span<const std::string_view> sortedUniquePrefixes);
+      const std::vector<std::string_view>& sortedUniquePrefixes);
 
   // returns 1 based indexes, 0 if not found.
   std::uint32_t findPrefix(std::string_view query) const noexcept;
@@ -298,8 +327,7 @@ LowerBoundPrefixMap<T>::LowerBoundPrefixMap(
       prefix2value.rend(),
       [](const auto& x, const auto& y) { return x.first == y.first; });
 
-  std::span<std::pair<std::string, T>> sortedUnique{
-      rend.base(), prefix2value.end()};
+  folly::Range sortedUnique{rend.base(), prefix2value.end()};
 
   std::vector<std::string_view> sortedPrefixes;
   sortedPrefixes.reserve(sortedUnique.size());

mcrouter/lib/fbi/cpp/test/LowerBoundPrefixMapTest.cpp

@@ -201,6 +201,8 @@ TEST(LowerBoundPrefixMapTest, OverrideValues) {
 TEST(LowerBoundPrefixMapTest, EmptyMap) {
   LowerBoundPrefixMap<int> lbMap;
   ASSERT_EQ(lbMap.findPrefix("a"), lbMap.end());
+  ASSERT_EQ(lbMap.size(), 0);
+  ASSERT_TRUE(lbMap.empty());
 }
 
 } // namespace

mcrouter/mcrouter_options_list.h

@@ -924,6 +924,14 @@ MCROUTER_OPTION_INTEGER(
     no_short,
     "1 in S non-error connection samples will be logged")
 
+MCROUTER_OPTION_TOGGLE(
+    enable_route_policy_v2,
+    false,
+    "enable-route-policy-v2",
+    no_short,
+    "Enable Route Policy V2")
+
+
 #ifdef ADDITIONAL_OPTIONS_FILE
 #include ADDITIONAL_OPTIONS_FILE
 #endif

mcrouter/routes/RootRoute.h

@@ -41,7 +41,8 @@ class RootRoute {
         rhMap_(
             routeSelectors,
             opts_.default_route,
-            opts_.send_invalid_route_to_default),
+            opts_.send_invalid_route_to_default,
+            opts_.enable_route_policy_v2),
         disableBroadcastDeleteRpc_(disableBroadcastDeleteRpc) {}
 
   template <class Request>

mcrouter/routes/RouteHandleMap-inl.h

@@ -56,12 +56,14 @@ using UniqueVectorMap = std::unordered_map<
 template <class RouteHandleIf>
 std::shared_ptr<RoutePolicyMap<RouteHandleIf>> makePolicyMap(
     UniqueVectorMap<RouteHandleIf>& uniqueVectors,
-    const RouteSelectorVector<RouteHandleIf>& v) {
+    const RouteSelectorVector<RouteHandleIf>& v,
+    bool enableRoutePolicyV2) {
   auto it = uniqueVectors.find(v);
   if (it != uniqueVectors.end()) {
     return it->second;
   }
-  return uniqueVectors[v] = std::make_shared<RoutePolicyMap<RouteHandleIf>>(v);
+  return uniqueVectors[v] = std::make_shared<RoutePolicyMap<RouteHandleIf>>(
+             v, enableRoutePolicyV2);
 }
 
 } // namespace detail
@@ -70,9 +72,11 @@ template <class RouteHandleIf>
 RouteHandleMap<RouteHandleIf>::RouteHandleMap(
     const RouteSelectorMap<RouteHandleIf>& routeSelectors,
     const RoutingPrefix& defaultRoute,
-    bool sendInvalidRouteToDefault)
+    bool sendInvalidRouteToDefault,
+    bool enableRoutePolicyV2)
     : defaultRoute_(defaultRoute),
-      sendInvalidRouteToDefault_(sendInvalidRouteToDefault) {
+      sendInvalidRouteToDefault_(sendInvalidRouteToDefault),
+      enableRoutePolicyV2_(enableRoutePolicyV2) {
   checkLogic(
       routeSelectors.find(defaultRoute_) != routeSelectors.end(),
       "invalid default route: {}",
@@ -111,12 +115,16 @@ RouteHandleMap<RouteHandleIf>::RouteHandleMap(
 
   // create corresponding RoutePolicyMaps
   detail::UniqueVectorMap<RouteHandleIf> uniqueVectors;
-  allRoutes_ = makePolicyMap(uniqueVectors, allRoutes);
+  allRoutes_ = makePolicyMap(uniqueVectors, allRoutes, enableRoutePolicyV2_);
   for (const auto& it : byRegion) {
-    byRegion_.emplace(it.first, makePolicyMap(uniqueVectors, it.second));
+    byRegion_.emplace(
+        it.first,
+        makePolicyMap(uniqueVectors, it.second, enableRoutePolicyV2_));
   }
   for (const auto& it : byRoute) {
-    byRoute_.emplace(it.first, makePolicyMap(uniqueVectors, it.second));
+    byRoute_.emplace(
+        it.first,
+        makePolicyMap(uniqueVectors, it.second, enableRoutePolicyV2_));
   }
 
   assert(byRoute_.find(defaultRoute_) != byRoute_.end());

mcrouter/routes/RouteHandleMap.h

@@ -34,7 +34,8 @@ class RouteHandleMap {
   RouteHandleMap(
       const RouteSelectorMap<RouteHandleIf>& routeSelectors,
       const RoutingPrefix& defaultRoute,
-      bool sendInvalidRouteToDefault);
+      bool sendInvalidRouteToDefault,
+      bool enableRoutePolicyV2);
 
   /**
    * @return pointer to a precalculated vector of route handles that a request
@@ -57,6 +58,7 @@ class RouteHandleMap {
   const std::vector<std::shared_ptr<RouteHandleIf>> emptyV_;
   const RoutingPrefix& defaultRoute_;
   bool sendInvalidRouteToDefault_;
+  bool enableRoutePolicyV2_;
   std::shared_ptr<RoutePolicyMap<RouteHandleIf>> defaultRouteMap_;
 
   std::shared_ptr<RoutePolicyMap<RouteHandleIf>> allRoutes_;

mcrouter/routes/RoutePolicyMap-inl.h

@@ -50,7 +50,13 @@ std::vector<std::shared_ptr<RouteHandleIf>> overrideItems(
 template <class RouteHandleIf>
 RoutePolicyMap<RouteHandleIf>::RoutePolicyMap(
     const std::vector<std::shared_ptr<PrefixSelectorRoute<RouteHandleIf>>>&
-        clusters) {
+        clusters,
+    bool useV2) {
+  if (useV2) {
+    v2_ = RoutePolicyMapV2<RouteHandleIf>(clusters);
+    return;
+  }
+
   // wildcards of all clusters
   std::vector<std::shared_ptr<RouteHandleIf>> wildcards;
   wildcards.reserve(clusters.size());
@@ -90,9 +96,84 @@ RoutePolicyMap<RouteHandleIf>::RoutePolicyMap(
 template <class RouteHandleIf>
 const std::vector<std::shared_ptr<RouteHandleIf>>&
 RoutePolicyMap<RouteHandleIf>::getTargetsForKey(folly::StringPiece key) const {
+  // empty means not initialized - i.e. the flag was not enabled.
+  if (!v2_.empty()) {
+    return v2_.getTargetsForKey(key);
+  }
   auto result = ut_.findPrefix(key);
   return result == ut_.end() ? emptyV_ : result->second;
 }
+
+template <class RouteHandleIf>
+RoutePolicyMapV2<RouteHandleIf>::RoutePolicyMapV2(
+    const std::vector<SharedClusterPtr>& clusters) {
+  typename LBRouteMap::Builder builder;
+
+  // not enough but it's not a problem, we will avoid at least some
+  // reallocations.
+  builder.reserve(clusters.size() + 1);
+  builder.insert({"", populateWildCards(clusters)});
+
+  // Combining routes for each key
+  folly::F14FastSet<std::string_view> seen;
+  for (const auto& cluster : clusters) {
+    for (const auto& [key, _] : cluster->policies) {
+      if (seen.insert(key).second) {
+        builder.insert({key, populateRoutesForKey(key, clusters)});
+      }
+    }
+  }
+
+  ut_ = std::move(builder).build();
+}
+
+template <class RouteHandleIf>
+// static
+auto RoutePolicyMapV2<RouteHandleIf>::populateWildCards(
+    const std::vector<SharedClusterPtr>& clusters)
+    -> std::vector<SharedRoutePtr> {
+  std::vector<SharedRoutePtr> res;
+  res.reserve(clusters.size());
+
+  folly::F14FastSet<const RouteHandleIf*> seen;
+  for (const auto& cluster : clusters) {
+    const SharedRoutePtr& ptr = cluster->wildcard;
+    if (ptr && seen.insert(ptr.get()).second) {
+      res.push_back(cluster->wildcard);
+    }
+  }
+
+  // This code is cold but the vector stays around for a while
+  res.shrink_to_fit();
+  return res;
+}
+
+template <class RouteHandleIf>
+// static
+auto RoutePolicyMapV2<RouteHandleIf>::populateRoutesForKey(
+    std::string_view key,
+    const std::vector<SharedClusterPtr>& clusters)
+    -> std::vector<SharedRoutePtr> {
+  std::vector<std::shared_ptr<RouteHandleIf>> res;
+  res.reserve(clusters.size());
+
+  folly::F14FastSet<const RouteHandleIf*> seen;
+  seen.reserve(clusters.size());
+
+  for (const auto& cluster : clusters) {
+    auto found = cluster->policies.findPrefix(key);
+    const SharedRoutePtr& ptr =
+        found == cluster->policies.end() ? cluster->wildcard : found->second;
+    if (ptr && seen.insert(ptr.get()).second) {
+      res.push_back(ptr);
+    }
+  }
+
+  // This code is cold but the vector stays around for a while
+  res.shrink_to_fit();
+  return res;
+}
+
 } // namespace mcrouter
 } // namespace memcache
 } // namespace facebook

mcrouter/routes/RoutePolicyMap.h

@@ -8,10 +8,13 @@
 #pragma once
 
 #include <memory>
+#include <string_view>
 #include <vector>
 
 #include <folly/Range.h>
+#include <folly/container/F14Set.h>
 
+#include "mcrouter/lib/fbi/cpp/LowerBoundPrefixMap.h"
 #include "mcrouter/lib/fbi/cpp/Trie.h"
 
 namespace facebook {
@@ -44,12 +47,49 @@ class PrefixSelectorRoute;
  * 2) use getTargetsForKey to get precalculated vector of targets.
  *    Complexity is O(min(longest key prefix in config, key length))
  */
+template <class RouteHandleIf>
+class RoutePolicyMapV2 {
+ public:
+  using SharedRoutePtr = std::shared_ptr<RouteHandleIf>;
+  using SharedClusterPtr = std::shared_ptr<PrefixSelectorRoute<RouteHandleIf>>;
+
+  RoutePolicyMapV2() = default;
+  explicit RoutePolicyMapV2(const std::vector<SharedClusterPtr>& clusters);
+
+  // NOTE: change to span when the migration is over
+  const std::vector<SharedRoutePtr>& getTargetsForKey(
+      std::string_view key) const {
+    // has to be there by construction
+    auto found = ut_.findPrefix(key);
+    // RoutePolicyMap always has wildcard which matches everything.
+    CHECK(found != ut_.end());
+    return found->value();
+  }
+
+  bool empty() const {
+    return ut_.empty();
+  }
+
+ private:
+  static std::vector<SharedRoutePtr> populateWildCards(
+      const std::vector<SharedClusterPtr>& clusters);
+
+  static std::vector<SharedRoutePtr> populateRoutesForKey(
+      std::string_view key,
+      const std::vector<SharedClusterPtr>& clusters);
+
+  using LBRouteMap = LowerBoundPrefixMap<std::vector<SharedRoutePtr>>;
+
+  LBRouteMap ut_;
+};
+
 template <class RouteHandleIf>
 class RoutePolicyMap {
  public:
   explicit RoutePolicyMap(
       const std::vector<std::shared_ptr<PrefixSelectorRoute<RouteHandleIf>>>&
-          clusters);
+          clusters,
+      bool useV2);
 
   /**
    * @return vector of route handles that a request with given key should be
@@ -59,6 +99,7 @@ class RoutePolicyMap {
       folly::StringPiece key) const;
 
  private:
+  RoutePolicyMapV2<RouteHandleIf> v2_;
   const std::vector<std::shared_ptr<RouteHandleIf>> emptyV_;
   /**
    * This Trie contains targets for each key prefix. It is built like this:
@@ -68,6 +109,7 @@ class RoutePolicyMap {
    */
   Trie<std::vector<std::shared_ptr<RouteHandleIf>>> ut_;
 };
+
 } // namespace mcrouter
 } // namespace memcache
 } // namespace facebook