local-rate-limit: improving cluster fetching validation (#41246)

Commit Message: local-rate-limit: improving cluster fetching validation

Signed-off-by: Adi Suissa-Peleg <[email protected]>

Author: adisuissa

envoy/upstream/cluster_manager.h

@@ -349,6 +349,16 @@ class ClusterManager {
    */
   virtual OptRef<const Cluster> getActiveCluster(const std::string& cluster_name) const PURE;
 
+  /**
+   * Receives a cluster name and returns an active or warming cluster (if found).
+   * @param cluster_name the name of the cluster.
+   * @return OptRef<const Cluster> A reference to the cluster if found, and nullopt otherwise.
+   *
+   * NOTE: This method is only thread safe on the main thread. It should not be called elsewhere.
+   */
+  virtual OptRef<const Cluster>
+  getActiveOrWarmingCluster(const std::string& cluster_name) const PURE;
+
   /**
    * Returns true iff the given cluster name is known in the cluster-manager
    * (either as active or as warming).

source/common/upstream/cluster_manager_impl.h

@@ -279,6 +279,17 @@ class ClusterManagerImpl : public ClusterManager,
     return absl::nullopt;
   }
 
+  OptRef<const Cluster> getActiveOrWarmingCluster(const std::string& cluster_name) const override {
+    ASSERT_IS_MAIN_OR_TEST_THREAD();
+    if (const auto& it = active_clusters_.find(cluster_name); it != active_clusters_.end()) {
+      return *it->second->cluster_;
+    }
+    if (const auto& it = warming_clusters_.find(cluster_name); it != warming_clusters_.end()) {
+      return *it->second->cluster_;
+    }
+    return absl::nullopt;
+  }
+
   bool hasCluster(const std::string& cluster_name) const override {
     ASSERT_IS_MAIN_OR_TEST_THREAD();
     return active_clusters_.contains(cluster_name) || warming_clusters_.contains(cluster_name);

source/extensions/filters/common/local_ratelimit/local_ratelimit_impl.cc

@@ -65,7 +65,7 @@ ShareProviderManagerSharedPtr ShareProviderManager::singleton(Event::Dispatcher&
         if (!local_cluster_name.has_value()) {
           return nullptr;
         }
-        auto cluster = cm.clusters().getCluster(local_cluster_name.value());
+        auto cluster = cm.getActiveOrWarmingCluster(local_cluster_name.value());
         if (!cluster.has_value()) {
           return nullptr;
         }

test/common/upstream/cluster_manager_impl_test.cc

@@ -1520,6 +1520,77 @@ TEST_F(ClusterManagerImplTest, OriginalDstInitialization) {
   factory_.tls_.shutdownThread();
 }
 
+TEST_F(ClusterManagerImplTest, GetActiveOrWarmingCluster) {
+  // Start with a static cluster.
+  const std::string bootstrap_yaml = R"EOF(
+  static_resources:
+    clusters:
+    - name: static_cluster
+      connect_timeout: 0.250s
+      type: static
+      lb_policy: round_robin
+      load_assignment:
+        cluster_name: static_cluster
+        endpoints:
+        - lb_endpoints:
+          - endpoint:
+              address:
+                socket_address:
+                  address: 127.0.0.1
+                  port_value: 11001
+  )EOF";
+  create(parseBootstrapFromV3Yaml(bootstrap_yaml));
+
+  // Static cluster should be active.
+  EXPECT_NE(absl::nullopt, cluster_manager_->getActiveCluster("static_cluster"));
+  EXPECT_NE(absl::nullopt, cluster_manager_->getActiveOrWarmingCluster("static_cluster"));
+  EXPECT_EQ(absl::nullopt, cluster_manager_->getActiveOrWarmingCluster("non_existent_cluster"));
+
+  // Now, add a dynamic cluster. It will start in warming state.
+  const std::string warming_cluster_yaml = R"EOF(
+    name: warming_cluster
+    connect_timeout: 0.250s
+    type: EDS
+    eds_cluster_config:
+      eds_config:
+        api_config_source:
+          api_type: GRPC
+          grpc_services:
+            envoy_grpc:
+              cluster_name: static_cluster
+  )EOF";
+  auto warming_cluster_config = parseClusterFromV3Yaml(warming_cluster_yaml);
+
+  // Mock the cluster creation for the warming cluster.
+  std::shared_ptr<MockClusterMockPrioritySet> warming_cluster =
+      std::make_shared<NiceMock<MockClusterMockPrioritySet>>();
+  warming_cluster->info_->name_ = "warming_cluster";
+  std::function<void()> cluster_init_callback;
+  EXPECT_CALL(*warming_cluster, initialize(_)).WillOnce(SaveArg<0>(&cluster_init_callback));
+  EXPECT_CALL(factory_, clusterFromProto_(ProtoEq(warming_cluster_config), _, true))
+      .WillOnce(Return(std::make_pair(warming_cluster, nullptr)));
+
+  // Add the cluster.
+  EXPECT_TRUE(*cluster_manager_->addOrUpdateCluster(warming_cluster_config, "version1"));
+
+  // The cluster should be in warming, not active.
+  EXPECT_EQ(absl::nullopt, cluster_manager_->getActiveCluster("warming_cluster"));
+  OptRef<const Cluster> cluster = cluster_manager_->getActiveOrWarmingCluster("warming_cluster");
+  EXPECT_NE(absl::nullopt, cluster);
+  EXPECT_EQ("warming_cluster", cluster->info()->name());
+
+  // Finish initialization. This should move it to active.
+  cluster_init_callback();
+
+  // Now the cluster should be active.
+  cluster = cluster_manager_->getActiveCluster("warming_cluster");
+  EXPECT_NE(absl::nullopt, cluster);
+  EXPECT_EQ("warming_cluster", cluster->info()->name());
+  cluster = cluster_manager_->getActiveOrWarmingCluster("warming_cluster");
+  EXPECT_NE(absl::nullopt, cluster);
+  EXPECT_EQ("warming_cluster", cluster->info()->name());
+}
+
 TEST_F(ClusterManagerImplTest, UpstreamSocketOptionsPassedToTcpConnPool) {
   createWithBasicStaticCluster();
   NiceMock<MockLoadBalancerContext> context;

test/mocks/upstream/cluster_manager.cc

@@ -38,17 +38,22 @@ MockClusterManager::MockClusterManager()
 MockClusterManager::~MockClusterManager() = default;
 
 void MockClusterManager::initializeClusters(const std::vector<std::string>& active_cluster_names,
-                                            const std::vector<std::string>&) {
+                                            const std::vector<std::string>& warming_cluster_names) {
   active_clusters_.clear();
+  warming_clusters_.clear();
   ClusterManager::ClusterInfoMaps info_map;
   for (const auto& name : active_cluster_names) {
     auto new_cluster = std::make_unique<NiceMock<MockCluster>>();
     new_cluster->info_->name_ = name;
     info_map.active_clusters_.emplace(name, *new_cluster);
     active_clusters_.emplace(name, std::move(new_cluster));
   }
-
-  // TODO(mattklein123): Add support for warming clusters when needed.
+  for (const auto& name : warming_cluster_names) {
+    auto new_cluster = std::make_unique<NiceMock<MockCluster>>();
+    new_cluster->info_->name_ = name;
+    info_map.warming_clusters_.emplace(name, *new_cluster);
+    warming_clusters_.emplace(name, std::move(new_cluster));
+  }
 
   ON_CALL(*this, clusters()).WillByDefault(Return(info_map));
   ON_CALL(*this, getActiveCluster(_))
@@ -58,9 +63,20 @@ void MockClusterManager::initializeClusters(const std::vector<std::string>& acti
         }
         return absl::nullopt;
       }));
+  ON_CALL(*this, getActiveOrWarmingCluster(_))
+      .WillByDefault(Invoke([this](const std::string& cluster_name) -> OptRef<const Cluster> {
+        if (const auto& it = active_clusters_.find(cluster_name); it != active_clusters_.end()) {
+          return *it->second;
+        }
+        if (const auto& it = warming_clusters_.find(cluster_name); it != warming_clusters_.end()) {
+          return *it->second;
+        }
+        return absl::nullopt;
+      }));
   ON_CALL(*this, hasCluster(_))
       .WillByDefault(Invoke([this](const std::string& cluster_name) -> bool {
-        return active_clusters_.find(cluster_name) != active_clusters_.end();
+        return active_clusters_.find(cluster_name) != active_clusters_.end() ||
+               warming_clusters_.find(cluster_name) != warming_clusters_.end();
       }));
   ON_CALL(*this, hasActiveClusters()).WillByDefault(Return(!active_cluster_names.empty()));
 }

test/mocks/upstream/cluster_manager.h

@@ -44,6 +44,8 @@ class MockClusterManager : public ClusterManager {
               (const envoy::config::bootstrap::v3::Bootstrap& bootstrap));
   MOCK_METHOD(ClusterInfoMaps, clusters, (), (const));
   MOCK_METHOD(OptRef<const Cluster>, getActiveCluster, (const std::string& cluster_name), (const));
+  MOCK_METHOD(OptRef<const Cluster>, getActiveOrWarmingCluster, (const std::string& cluster_name),
+              (const));
   MOCK_METHOD(bool, hasCluster, (const std::string& cluster_name), (const));
   MOCK_METHOD(bool, hasActiveClusters, (), (const));