http ratelimit: option to reduce budget on stream done

api/envoy/config/route/v3/route_components.proto

@@ -1868,6 +1868,7 @@ message VirtualCluster {
 
 // Global rate limiting :ref:`architecture overview <arch_overview_global_rate_limit>`.
 // Also applies to Local rate limiting :ref:`using descriptors <config_http_filters_local_rate_limit_descriptors>`.
+// [#next-free-field: 6]
 message RateLimit {
   option (udpa.annotations.versioning).previous_message_type = "envoy.api.v2.route.RateLimit";
 
@@ -2193,6 +2194,23 @@ message RateLimit {
   // from metadata, no override is provided. See :ref:`rate limit override
   // <config_http_filters_rate_limit_rate_limit_override>` for more information.
   Override limit = 4;
+
+  // If true, the rate limit request will be applied when the stream completes. The default value is false.
+  // This is useful when the rate limit budget needs to reflect the response context that is not available
+  // on the request path.
+  //
+  // For example, let's say the upstream service calculates the usage statistics and returns them in the response body
+  // and we want to utilize these numbers to apply the rate limit action for the subsequent requests.
+  // Combined with another filter that can set the desired addend based on the response (e.g. Lua filter),
+  // this can be used to subtract the usage statistics from the rate limit budget.
+  //
+  // A rate limit applied on the stream completion is "fire-and-forget" by nature, and rate limit is not enforced by this config.
+  // In other words, the current request won't be blocked when this is true, but the budget will be updated for the subsequent
+  // requests based on the action with this field set to true. Users should ensure that the rate limit is enforced by the actions
+  // applied on the request path, i.e. the ones with this field set to false.
+  //
+  // Currently, this is only supported by the HTTP global rate filter.
+  bool apply_on_stream_done = 5;
 }
 
 // .. attention::

envoy/router/router_ratelimit.h

@@ -48,6 +48,11 @@ class RateLimitPolicyEntry {
    */
   virtual const std::string& disableKey() const PURE;
 
+  /**
+   * @return true if this rate limit policy should be applied on stream done.
+   */
+  virtual bool applyOnStreamDone() const PURE;
+
   /**
    * Potentially populate the descriptor array with new descriptors to query.
    * @param descriptors supplies the descriptor array to optionally fill.

source/common/router/router_ratelimit.cc

@@ -267,7 +267,8 @@ RateLimitPolicyEntryImpl::RateLimitPolicyEntryImpl(
     const envoy::config::route::v3::RateLimit& config,
     Server::Configuration::CommonFactoryContext& context, absl::Status& creation_status)
     : disable_key_(config.disable_key()),
-      stage_(static_cast<uint64_t>(PROTOBUF_GET_WRAPPED_OR_DEFAULT(config, stage, 0))) {
+      stage_(static_cast<uint64_t>(PROTOBUF_GET_WRAPPED_OR_DEFAULT(config, stage, 0))),
+      apply_on_stream_done_(config.apply_on_stream_done()) {
   for (const auto& action : config.actions()) {
     switch (action.action_specifier_case()) {
     case envoy::config::route::v3::RateLimit::Action::ActionSpecifierCase::kSourceCluster:

source/common/router/router_ratelimit.h

@@ -256,12 +256,14 @@ class RateLimitPolicyEntryImpl : public RateLimitPolicyEntry {
                                 const std::string& local_service_cluster,
                                 const Http::RequestHeaderMap&,
                                 const StreamInfo::StreamInfo& info) const override;
+  bool applyOnStreamDone() const override { return apply_on_stream_done_; }
 
 private:
   const std::string disable_key_;
   uint64_t stage_;
   std::vector<RateLimit::DescriptorProducerPtr> actions_;
   absl::optional<RateLimitOverrideActionPtr> limit_override_ = absl::nullopt;
+  bool apply_on_stream_done_;
 };
 
 /**

source/extensions/filters/common/ratelimit/ratelimit.h

@@ -90,7 +90,7 @@ class Client {
    */
   virtual void limit(RequestCallbacks& callbacks, const std::string& domain,
                      const std::vector<Envoy::RateLimit::Descriptor>& descriptors,
-                     Tracing::Span& parent_span, const StreamInfo::StreamInfo& stream_info,
+                     Tracing::Span& parent_span, OptRef<const StreamInfo::StreamInfo> stream_info,
                      uint32_t hits_addend) PURE;
 };
 

source/extensions/filters/common/ratelimit/ratelimit_impl.cc

@@ -59,18 +59,19 @@ void GrpcClientImpl::createRequest(envoy::service::ratelimit::v3::RateLimitReque
 
 void GrpcClientImpl::limit(RequestCallbacks& callbacks, const std::string& domain,
                            const std::vector<Envoy::RateLimit::Descriptor>& descriptors,
-                           Tracing::Span& parent_span, const StreamInfo::StreamInfo& stream_info,
-                           uint32_t hits_addend) {
+                           Tracing::Span& parent_span,
+                           OptRef<const StreamInfo::StreamInfo> stream_info, uint32_t hits_addend) {
   ASSERT(callbacks_ == nullptr);
   callbacks_ = &callbacks;
 
   envoy::service::ratelimit::v3::RateLimitRequest request;
   createRequest(request, domain, descriptors, hits_addend);
 
-  request_ =
-      async_client_->send(service_method_, request, *this, parent_span,
-                          Http::AsyncClient::RequestOptions().setTimeout(timeout_).setParentContext(
-                              Http::AsyncClient::ParentContext{&stream_info}));
+  auto options = Http::AsyncClient::RequestOptions().setTimeout(timeout_);
+  if (stream_info) {
+    options.setParentContext(Http::AsyncClient::ParentContext{&*stream_info});
+  }
+  request_ = async_client_->send(service_method_, request, *this, parent_span, options);
 }
 
 void GrpcClientImpl::onSuccess(

source/extensions/filters/common/ratelimit/ratelimit_impl.h

@@ -57,7 +57,7 @@ class GrpcClientImpl : public Client,
   void cancel() override;
   void limit(RequestCallbacks& callbacks, const std::string& domain,
              const std::vector<Envoy::RateLimit::Descriptor>& descriptors,
-             Tracing::Span& parent_span, const StreamInfo::StreamInfo& stream_info,
+             Tracing::Span& parent_span, OptRef<const StreamInfo::StreamInfo> stream_info,
              uint32_t hits_addend = 0) override;
 
   // Grpc::AsyncRequestCallbacks

source/extensions/filters/http/ratelimit/ratelimit.cc

@@ -1,5 +1,6 @@
 #include "source/extensions/filters/http/ratelimit/ratelimit.h"
 
+#include <memory>
 #include <string>
 #include <vector>
 
@@ -86,21 +87,24 @@ void Filter::initiateCall(const Http::RequestHeaderMap& headers) {
     break;
   }
 
+  if (!descriptors.empty()) {
+    state_ = State::Calling;
+    initiating_call_ = true;
+    client_->limit(*this, getDomain(), descriptors, callbacks_->activeSpan(),
+                   callbacks_->streamInfo(), getHitAddend());
+    initiating_call_ = false;
+  }
+}
+
+double Filter::getHitAddend() {
   const StreamInfo::UInt32Accessor* hits_addend_filter_state =
       callbacks_->streamInfo().filterState()->getDataReadOnly<StreamInfo::UInt32Accessor>(
           HitsAddendFilterStateKey);
   double hits_addend = 0;
   if (hits_addend_filter_state != nullptr) {
     hits_addend = hits_addend_filter_state->value();
   }
-
-  if (!descriptors.empty()) {
-    state_ = State::Calling;
-    initiating_call_ = true;
-    client_->limit(*this, getDomain(), descriptors, callbacks_->activeSpan(),
-                   callbacks_->streamInfo(), hits_addend);
-    initiating_call_ = false;
-  }
+  return hits_addend;
 }
 
 Http::FilterHeadersStatus Filter::decodeHeaders(Http::RequestHeaderMap& headers, bool) {
@@ -161,6 +165,18 @@ void Filter::onDestroy() {
   if (state_ == State::Calling) {
     state_ = State::Complete;
     client_->cancel();
+  } else {
+    // If the filter doesn't have an outstanding limit request (made during decodeHeaders) and has
+    // descriptors, then we can apply the rate limit on stream done if the config allows it.
+    if (!descriptors_apply_on_stream_done_.empty()) {
+      client_->cancel(); // Clears the internal state of the client, so that we can reuse it.
+      // Since this filter is being destroyed, we need to keep the client alive until the request
+      // is complete.
+      auto callback = new OnStreamDoneCallBack(std::move(client_));
+      auto& ref = *callback;
+      callback->client()->limit(ref, getDomain(), descriptors_apply_on_stream_done_,
+                                Tracing::NullSpan::instance(), absl::nullopt, getHitAddend());
+    }
   }
 }
 
@@ -258,7 +274,7 @@ void Filter::complete(Filters::Common::RateLimit::LimitStatus status,
 
 void Filter::populateRateLimitDescriptors(const Router::RateLimitPolicy& rate_limit_policy,
                                           std::vector<RateLimit::Descriptor>& descriptors,
-                                          const Http::RequestHeaderMap& headers) const {
+                                          const Http::RequestHeaderMap& headers) {
   for (const Router::RateLimitPolicyEntry& rate_limit :
        rate_limit_policy.getApplicableRateLimit(config_->stage())) {
     const std::string& disable_key = rate_limit.disableKey();
@@ -267,8 +283,16 @@ void Filter::populateRateLimitDescriptors(const Router::RateLimitPolicy& rate_li
             fmt::format("ratelimit.{}.http_filter_enabled", disable_key), 100)) {
       continue;
     }
-    rate_limit.populateDescriptors(descriptors, config_->localInfo().clusterName(), headers,
-                                   callbacks_->streamInfo());
+    if (rate_limit.applyOnStreamDone()) {
+      // If the rate limit policy is to be applied on stream done, we populate the descriptors in a
+      // separate vector. This vector will be used to apply the rate limit on stream done.
+      rate_limit.populateDescriptors(descriptors_apply_on_stream_done_,
+                                     config_->localInfo().clusterName(), headers,
+                                     callbacks_->streamInfo());
+    } else {
+      rate_limit.populateDescriptors(descriptors, config_->localInfo().clusterName(), headers,
+                                     callbacks_->streamInfo());
+    }
   }
 }
 
@@ -330,6 +354,14 @@ std::string Filter::getDomain() {
   return config_->domain();
 }
 
+void OnStreamDoneCallBack::complete(Filters::Common::RateLimit::LimitStatus,
+                                    Filters::Common::RateLimit::DescriptorStatusListPtr&&,
+                                    Http::ResponseHeaderMapPtr&&, Http::RequestHeaderMapPtr&&,
+                                    const std::string&,
+                                    Filters::Common::RateLimit::DynamicMetadataPtr&&) {
+  delete this;
+}
+
 } // namespace RateLimitFilter
 } // namespace HttpFilters
 } // namespace Extensions

source/extensions/filters/http/ratelimit/ratelimit.h

@@ -187,9 +187,10 @@ class Filter : public Http::StreamFilter, public Filters::Common::RateLimit::Req
   void initiateCall(const Http::RequestHeaderMap& headers);
   void populateRateLimitDescriptors(const Router::RateLimitPolicy& rate_limit_policy,
                                     std::vector<Envoy::RateLimit::Descriptor>& descriptors,
-                                    const Http::RequestHeaderMap& headers) const;
+                                    const Http::RequestHeaderMap& headers);
   void populateResponseHeaders(Http::HeaderMap& response_headers, bool from_local_reply);
   void appendRequestHeaders(Http::HeaderMapPtr& request_headers_to_add);
+  double getHitAddend();
   VhRateLimitOptions getVirtualHostRateLimitOption(const Router::RouteConstSharedPtr& route);
   std::string getDomain();
 
@@ -206,6 +207,30 @@ class Filter : public Http::StreamFilter, public Filters::Common::RateLimit::Req
   bool initiating_call_{};
   Http::ResponseHeaderMapPtr response_headers_to_add_;
   Http::RequestHeaderMap* request_headers_{};
+  // Holds the descriptors that should be applied on stream done which will be populated during
+  // decodeHeaders.
+  std::vector<Envoy::RateLimit::Descriptor> descriptors_apply_on_stream_done_{};
+};
+
+/**
+ * This implements the rate limit callback that outlives the filter holding the client.
+ * On completion, it deletes itself.
+ */
+class OnStreamDoneCallBack : public Filters::Common::RateLimit::RequestCallbacks {
+public:
+  OnStreamDoneCallBack(Filters::Common::RateLimit::ClientPtr client) : client_(std::move(client)) {}
+  ~OnStreamDoneCallBack() override = default;
+
+  // RateLimit::RequestCallbacks
+  void complete(Filters::Common::RateLimit::LimitStatus,
+                Filters::Common::RateLimit::DescriptorStatusListPtr&&, Http::ResponseHeaderMapPtr&&,
+                Http::RequestHeaderMapPtr&&, const std::string&,
+                Filters::Common::RateLimit::DynamicMetadataPtr&&) override;
+
+  Filters::Common::RateLimit::ClientPtr& client() { return client_; }
+
+private:
+  Filters::Common::RateLimit::ClientPtr client_;
 };
 
 } // namespace RateLimitFilter

test/extensions/filters/common/ratelimit/mocks.h

@@ -26,7 +26,7 @@ class MockClient : public Client {
   MOCK_METHOD(void, limit,
               (RequestCallbacks & callbacks, const std::string& domain,
                const std::vector<Envoy::RateLimit::Descriptor>& descriptors,
-               Tracing::Span& parent_span, const StreamInfo::StreamInfo& stream_info,
+               Tracing::Span& parent_span, OptRef<const StreamInfo::StreamInfo> stream_info,
                uint32_t hits_addend));
 };
 

test/extensions/filters/http/ratelimit/ratelimit_test.cc

@@ -265,7 +265,7 @@ TEST_F(HttpRateLimitFilterTest, OkResponseWithAdditionalHitsAddend) {
 
   filter_callbacks_.stream_info_.filter_state_->setData(
       "envoy.ratelimit.hits_addend", std::make_unique<StreamInfo::UInt32AccessorImpl>(5),
-      StreamInfo::FilterState::StateType::ReadOnly);
+      StreamInfo::FilterState::StateType::Mutable);
   EXPECT_CALL(filter_callbacks_.route_->route_entry_.rate_limit_policy_, getApplicableRateLimit(0));
 
   EXPECT_CALL(route_rate_limit_, populateDescriptors(_, _, _, _))
@@ -274,6 +274,11 @@ TEST_F(HttpRateLimitFilterTest, OkResponseWithAdditionalHitsAddend) {
   EXPECT_CALL(filter_callbacks_.route_->virtual_host_.rate_limit_policy_,
               getApplicableRateLimit(0));
 
+  // Make sure that descriptor_two_ will be applied on stream done.
+  EXPECT_CALL(vh_rate_limit_, applyOnStreamDone()).WillOnce(Return(true));
+  EXPECT_CALL(vh_rate_limit_, populateDescriptors(_, _, _, _))
+      .WillOnce(SetArgReferee<0>(descriptor_two_));
+
   EXPECT_CALL(*client_, limit(_, "foo",
                               testing::ContainerEq(std::vector<RateLimit::Descriptor>{
                                   {{{"descriptor_key", "descriptor_value"}}}}),
@@ -304,6 +309,22 @@ TEST_F(HttpRateLimitFilterTest, OkResponseWithAdditionalHitsAddend) {
 
   EXPECT_EQ(
       1U, filter_callbacks_.clusterInfo()->statsScope().counterFromStatName(ratelimit_ok_).value());
+
+  // Test the behavior when apply_on_stream_done is true.
+  testing::Mock::VerifyAndClearExpectations(client_);
+  filter_callbacks_.stream_info_.filter_state_->setData(
+      // Ensures that addend can be set differently than the request path.
+      "envoy.ratelimit.hits_addend", std::make_unique<StreamInfo::UInt32AccessorImpl>(100),
+      StreamInfo::FilterState::StateType::Mutable);
+  EXPECT_CALL(*client_, cancel());
+  EXPECT_CALL(*client_, limit(_, "foo", testing::ContainerEq(descriptor_two_), _, _, 100))
+      .WillOnce(
+          WithArgs<0>(Invoke([&](Filters::Common::RateLimit::RequestCallbacks& callbacks) -> void {
+            request_callbacks_ = &callbacks;
+          })));
+  filter_->onDestroy();
+  request_callbacks_->complete(Filters::Common::RateLimit::LimitStatus::OK, nullptr, nullptr,
+                               nullptr, "", nullptr);
 }
 
 TEST_F(HttpRateLimitFilterTest, OkResponseWithHeaders) {

test/mocks/router/mocks.h

@@ -250,6 +250,7 @@ class MockRateLimitPolicyEntry : public RateLimitPolicyEntry {
                const std::string& local_service_cluster, const Http::RequestHeaderMap& headers,
                const StreamInfo::StreamInfo& info),
               (const));
+  MOCK_METHOD(bool, applyOnStreamDone, (), (const));
 
   uint64_t stage_{};
   std::string disable_key_;