http ratelimit: option to reduce budget on stream done

api/envoy/extensions/filters/http/ratelimit/v3/rate_limit.proto

@@ -25,7 +25,7 @@ option (udpa.annotations.file_status).package_version_status = ACTIVE;
 // Rate limit :ref:`configuration overview <config_http_filters_rate_limit>`.
 // [#extension: envoy.filters.http.ratelimit]
 
-// [#next-free-field: 14]
+// [#next-free-field: 15]
 message RateLimit {
   option (udpa.annotations.versioning).previous_message_type =
       "envoy.config.filter.http.rate_limit.v2.RateLimit";
@@ -134,6 +134,21 @@ message RateLimit {
   // Optional additional prefix to use when emitting statistics. This allows to distinguish
   // emitted statistics between configured ``ratelimit`` filters in an HTTP filter chain.
   string stat_prefix = 13;
+
+  // If true, rate limit requests will also be sent to the rate limit service when the stream completes.
+  // 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, 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.
+  //
+  // The rate limit requests sent on the stream completion are "fire-and-forget" by nature, and rate limit is not enforced
+  // on the current HTTP stream being completed. The filter will only update the budget for the subsequent requests at
+  // that point. Hence the effect of the rate limit requests made during the stream completion is not visible in the current
+  // but only in the subsequent requests.
+  bool apply_on_stream_done = 14;
 }
 
 // Global rate limiting :ref:`architecture overview <arch_overview_global_rate_limit>`.

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

@@ -65,39 +65,41 @@ void Filter::initiateCall(const Http::RequestHeaderMap& headers) {
     return;
   }
 
-  std::vector<Envoy::RateLimit::Descriptor> descriptors;
-
   const Router::RouteEntry* route_entry = route->routeEntry();
   // Get all applicable rate limit policy entries for the route.
-  populateRateLimitDescriptors(route_entry->rateLimitPolicy(), descriptors, headers);
+  populateRateLimitDescriptors(route_entry->rateLimitPolicy(), descriptors_, headers);
 
   VhRateLimitOptions vh_rate_limit_option = getVirtualHostRateLimitOption(route);
 
   switch (vh_rate_limit_option) {
   case VhRateLimitOptions::Ignore:
     break;
   case VhRateLimitOptions::Include:
-    populateRateLimitDescriptors(route->virtualHost().rateLimitPolicy(), descriptors, headers);
+    populateRateLimitDescriptors(route->virtualHost().rateLimitPolicy(), descriptors_, headers);
     break;
   case VhRateLimitOptions::Override:
     if (route_entry->rateLimitPolicy().empty()) {
-      populateRateLimitDescriptors(route->virtualHost().rateLimitPolicy(), descriptors, headers);
+      populateRateLimitDescriptors(route->virtualHost().rateLimitPolicy(), descriptors_, headers);
     }
     break;
   }
+  makeRateLimitRequest();
+}
 
-  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();
-  }
+void Filter::makeRateLimitRequest() {
+  if (!descriptors_.empty()) {
+    // TODO: Make addend configirable via substituion command.
+    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(),
+    client_->limit(*this, getDomain(), descriptors_, callbacks_->activeSpan(),
                    callbacks_->streamInfo(), hits_addend);
     initiating_call_ = false;
   }
@@ -158,9 +160,14 @@ Http::FilterMetadataStatus Filter::encodeMetadata(Http::MetadataMap&) {
 void Filter::setEncoderFilterCallbacks(Http::StreamEncoderFilterCallbacks&) {}
 
 void Filter::onDestroy() {
-  if (state_ == State::Calling) {
-    state_ = State::Complete;
-    client_->cancel();
+  if (config_->applyOnStreamDone()) {
+    state_ = State::OnStreamDone;
+    makeRateLimitRequest();
+  } else {
+    if (state_ == State::Calling) {
+      state_ = State::Complete;
+      client_->cancel();
+    }
   }
 }
 
@@ -170,6 +177,11 @@ void Filter::complete(Filters::Common::RateLimit::LimitStatus status,
                       Http::RequestHeaderMapPtr&& request_headers_to_add,
                       const std::string& response_body,
                       Filters::Common::RateLimit::DynamicMetadataPtr&& dynamic_metadata) {
+  if (state_ == State::OnStreamDone) {
+    // We have no more work to do as the rate limit request made during on completion is
+    // fire-and-forget.
+    return;
+  }
   state_ = State::Complete;
   response_headers_to_add_ = std::move(response_headers_to_add);
   Http::HeaderMapPtr req_headers_to_add = std::move(request_headers_to_add);

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

@@ -61,7 +61,8 @@ class FilterConfig {
         response_headers_parser_(THROW_OR_RETURN_VALUE(
             Envoy::Router::HeaderParser::configure(config.response_headers_to_add()),
             Router::HeaderParserPtr)),
-        status_on_error_(toRatelimitServerErrorCode(config.status_on_error().code())) {}
+        status_on_error_(toRatelimitServerErrorCode(config.status_on_error().code())),
+        apply_on_stream_done_(config.apply_on_stream_done()) {}
   const std::string& domain() const { return domain_; }
   const LocalInfo::LocalInfo& localInfo() const { return local_info_; }
   uint64_t stage() const { return stage_; }
@@ -79,6 +80,7 @@ class FilterConfig {
   Http::Code rateLimitedStatus() { return rate_limited_status_; }
   const Router::HeaderParser& responseHeadersParser() const { return *response_headers_parser_; }
   Http::Code statusOnError() const { return status_on_error_; }
+  bool applyOnStreamDone() const { return apply_on_stream_done_; }
 
 private:
   static FilterRequestType stringToType(const std::string& request_type) {
@@ -123,6 +125,7 @@ class FilterConfig {
   const Http::Code rate_limited_status_;
   Router::HeaderParserPtr response_headers_parser_;
   const Http::Code status_on_error_;
+  const bool apply_on_stream_done_;
 };
 
 using FilterConfigSharedPtr = std::shared_ptr<FilterConfig>;
@@ -185,6 +188,7 @@ class Filter : public Http::StreamFilter, public Filters::Common::RateLimit::Req
 
 private:
   void initiateCall(const Http::RequestHeaderMap& headers);
+  void makeRateLimitRequest();
   void populateRateLimitDescriptors(const Router::RateLimitPolicy& rate_limit_policy,
                                     std::vector<Envoy::RateLimit::Descriptor>& descriptors,
                                     const Http::RequestHeaderMap& headers) const;
@@ -195,7 +199,7 @@ class Filter : public Http::StreamFilter, public Filters::Common::RateLimit::Req
 
   Http::Context& httpContext() { return config_->httpContext(); }
 
-  enum class State { NotStarted, Calling, Complete, Responded };
+  enum class State { NotStarted, Calling, Complete, Responded, OnStreamDone };
 
   FilterConfigSharedPtr config_;
   Filters::Common::RateLimit::ClientPtr client_;
@@ -206,6 +210,7 @@ class Filter : public Http::StreamFilter, public Filters::Common::RateLimit::Req
   bool initiating_call_{};
   Http::ResponseHeaderMapPtr response_headers_to_add_;
   Http::RequestHeaderMap* request_headers_{};
+  std::vector<Envoy::RateLimit::Descriptor> descriptors_{};
 };
 
 } // namespace RateLimitFilter