ratelimit: add support for x-ratelimit-* headers in local rate limiting #18157
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,5 +1,7 @@ | ||
| #include "source/extensions/filters/common/local_ratelimit/local_ratelimit_impl.h" | ||
|
|
||
| #include <chrono> | ||
|
|
||
| #include "source/common/protobuf/utility.h" | ||
|
|
||
| namespace Envoy { | ||
|
|
@@ -21,10 +23,15 @@ LocalRateLimiterImpl::LocalRateLimiterImpl( | |
| throw EnvoyException("local rate limit token bucket fill timer must be >= 50ms"); | ||
| } | ||
|
|
||
| RateLimit::TokenBucket default_token_bucket; | ||
| default_token_bucket.max_tokens_ = max_tokens; | ||
| default_token_bucket.tokens_per_fill_ = tokens_per_fill; | ||
| default_token_bucket.fill_interval_ = absl::FromChrono(fill_interval); | ||
|
Member
There was a problem hiding this comment. Nit: This
Contributor
Author
There was a problem hiding this comment. Thanks, I will fix it. |
||
| default_descriptor_.token_bucket_ = default_token_bucket; | ||
| auto default_token_state = std::make_unique<TokenState>(); | ||
| default_token_state->tokens_ = max_tokens; | ||
| default_token_state->fill_time_ = time_source_.monotonicTime(); | ||
|
Member
There was a problem hiding this comment. Github prevent me from commenting on that Fundamentally, it is the issue in network filter test not the code here: In the test Have you tried to see if it works by switching default testSystem in class If that doesn't work, maybe we can pull in folks who have more experience on time system.
Contributor
Author
There was a problem hiding this comment. I have performed a quick test by switching the time system in In the README, a So, shall we see and wait the change of the time system, or use the hacky way (in my commit) to switch the time system.
Member
There was a problem hiding this comment. The test results looks promising and the change is pretty self-contained in this test. I feel it might depends on How long/hard it will take to switch default to I am not an expert on time system, let me pull in @jmarantz who filed the TODO #4160 to shed some lights here.
Contributor
Author
There was a problem hiding this comment. I found it is not easy to switch to
Contributor
There was a problem hiding this comment. Sorry I missed this. in general we should only use one time system at a time, either a Simulated one or a Real one. The "GlobalTimeSystem" is just a helper class to get access to a time-system reference to the Real or Simulated system that already has been declared (or to declare a default real-time system if one is not already declared). To make a particular test use SimulatedTime you should only need to make its test-class derive from Event::TestUsingSimulatedTime, from test/test_common/simulated_time_system.h Now -- you may experience issues with simulated time if there are parts of the test that depend on real time and don't provide a mechanism to inject time. One place I think this might come up with is gRPC which doesn't (afaik) provide a time-injection hook, but may use sleeps and real-time timeouts that don't work well with simulated time.
Contributor
Author
There was a problem hiding this comment. Thanks for your reply. But we have a fuzz test which is also influenced by the time system. We cannot make the test class derive from |
||
| default_descriptor_.token_state_ = std::move(default_token_state); | ||
|
Member
There was a problem hiding this comment. IMO, Adding this new I see that you are using it for Please see my comment on that function for statusOr
Contributor
Author
There was a problem hiding this comment. It is great! I will take your advice and change the code. Thanks. |
||
|
|
||
| if (fill_timer_) { | ||
| fill_timer_->enableTimer(fill_interval); | ||
|
|
@@ -38,7 +45,8 @@ LocalRateLimiterImpl::LocalRateLimiterImpl( | |
| RateLimit::TokenBucket token_bucket; | ||
| token_bucket.fill_interval_ = | ||
| absl::Milliseconds(PROTOBUF_GET_MS_OR_DEFAULT(descriptor.token_bucket(), fill_interval, 0)); | ||
| if (token_bucket.fill_interval_ % default_descriptor_.token_bucket_.fill_interval_ != | ||
| absl::ZeroDuration()) { | ||
| throw EnvoyException( | ||
| "local rate descriptor limit is not a multiple of token bucket fill timer"); | ||
| } | ||
|
|
@@ -67,12 +75,13 @@ LocalRateLimiterImpl::~LocalRateLimiterImpl() { | |
| } | ||
|
|
||
| void LocalRateLimiterImpl::onFillTimer() { | ||
| onFillTimerHelper(*default_descriptor_.token_state_, default_descriptor_.token_bucket_); | ||
| onFillTimerDescriptorHelper(); | ||
| fill_timer_->enableTimer( | ||
| absl::ToChronoMilliseconds(default_descriptor_.token_bucket_.fill_interval_)); | ||
| } | ||
|
|
||
| void LocalRateLimiterImpl::onFillTimerHelper(TokenState& tokens, | ||
| const RateLimit::TokenBucket& bucket) { | ||
| // Relaxed consistency is used for all operations because we don't care about ordering, just the | ||
| // final atomic correctness. | ||
|
|
@@ -88,6 +97,9 @@ void LocalRateLimiterImpl::onFillTimerHelper(const TokenState& tokens, | |
| // Loop while the weak CAS fails trying to update the tokens value. | ||
| } while (!tokens.tokens_.compare_exchange_weak(expected_tokens, new_tokens_value, | ||
| std::memory_order_relaxed)); | ||
|
|
||
| // Update fill time at last. | ||
| tokens.fill_time_ = time_source_.monotonicTime(); | ||
| } | ||
|
|
||
| void LocalRateLimiterImpl::onFillTimerDescriptorHelper() { | ||
|
|
@@ -97,7 +109,6 @@ void LocalRateLimiterImpl::onFillTimerDescriptorHelper() { | |
| current_time - descriptor.token_state_->fill_time_) >= | ||
| absl::ToChronoMilliseconds(descriptor.token_bucket_.fill_interval_)) { | ||
| onFillTimerHelper(*descriptor.token_state_, descriptor.token_bucket_); | ||
| } | ||
| } | ||
| } | ||
|
|
@@ -123,17 +134,49 @@ bool LocalRateLimiterImpl::requestAllowedHelper(const TokenState& tokens) const | |
| return true; | ||
| } | ||
|
|
||
| const LocalRateLimiterImpl::LocalDescriptorImpl& LocalRateLimiterImpl::descriptorHelper( | ||
|
Member
There was a problem hiding this comment. You can refactor the code as below, which better reflect the intents of the code here. or raw pointer also works : |
||
| absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const { | ||
| if (!descriptors_.empty() && !request_descriptors.empty()) { | ||
| // The override rate limit descriptor is selected by the first full match from the request | ||
| // descriptors. | ||
| for (const auto& request_descriptor : request_descriptors) { | ||
| auto it = descriptors_.find(request_descriptor); | ||
| if (it != descriptors_.end()) { | ||
| return *it; | ||
| } | ||
| } | ||
| } | ||
| return default_descriptor_; | ||
| } | ||
|
|
||
| bool LocalRateLimiterImpl::requestAllowed( | ||
| absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const { | ||
| return requestAllowedHelper(*descriptorHelper(request_descriptors).token_state_); | ||
|
Member
There was a problem hiding this comment. With statusOr style, you can do something like |
||
| } | ||
|
|
||
| uint32_t LocalRateLimiterImpl::maxTokens( | ||
|
Contributor
There was a problem hiding this comment. Per the previous discussion on this maybe we should have these functions be called
Contributor
Author
There was a problem hiding this comment. I am not sure if we need this change because it would make the method names very long.
Contributor
There was a problem hiding this comment. We tend to favor longer names if it helps readability, but it's probably ok as is |
||
| absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const { | ||
| return descriptorHelper(request_descriptors).token_bucket_.max_tokens_; | ||
| } | ||
|
Comment on lines
+146
to
+160
Contributor
There was a problem hiding this comment. Should this be finding the max of all the descriptors? Seems like it's finding the max of the first descriptor we know about? Docs might be good here if this is the intended implementation, I'm not that familiar with this code and seeing the impl for a function called
Member
There was a problem hiding this comment. +1, it seems that we should return the max token of all matched request descriptors?
Contributor
Author
There was a problem hiding this comment. No, each request will correspond to a certain descriptor, and these descriptors have their own token bucket. So what we need to do here is to find the corresponding token bucket and return its
Member
There was a problem hiding this comment. Thanks for the explanation. Just for my own learning purpose: My previous comment about max_token was referring to such case. But I could be just wrong as I am not that familiar with code here.
Member
There was a problem hiding this comment. NVM. I found the answer:) It might be a good idea to add this comment to the code when you move this pattern into a helper method? in case other people are wondering this as well.
Contributor
Author
There was a problem hiding this comment. Thanks, I will add the comment. |
||
|
|
||
| uint32_t LocalRateLimiterImpl::remainingTokens( | ||
| absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const { | ||
| return descriptorHelper(request_descriptors) | ||
| .token_state_->tokens_.load(std::memory_order_relaxed); | ||
| } | ||
|
|
||
| int64_t LocalRateLimiterImpl::remainingFillInterval( | ||
| absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const { | ||
| using namespace std::literals; | ||
|
|
||
| auto current_time = time_source_.monotonicTime(); | ||
| auto& descriptor = descriptorHelper(request_descriptors); | ||
| ASSERT(std::chrono::duration_cast<std::chrono::milliseconds>( | ||
| current_time - descriptor.token_state_->fill_time_) <= | ||
| absl::ToChronoMilliseconds(descriptor.token_bucket_.fill_interval_)); | ||
| return absl::ToInt64Seconds( | ||
| descriptor.token_bucket_.fill_interval_ - | ||
| absl::Seconds((current_time - descriptor.token_state_->fill_time_) / 1s)); | ||
| } | ||
|
|
||
| } // namespace LocalRateLimit | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,28 @@ | ||
| #pragma once | ||
|
|
||
| #include "envoy/http/header_map.h" | ||
|
|
||
| #include "source/common/singleton/const_singleton.h" | ||
|
|
||
| namespace Envoy { | ||
| namespace Extensions { | ||
| namespace HttpFilters { | ||
| namespace Common { | ||
| namespace RateLimit { | ||
|
Member
There was a problem hiding this comment. nit: newline after namespace (would have thought clang-tidy would complain here).
Contributor
Author
There was a problem hiding this comment. Thanks. I will run clang-tidy and fix related issues. |
||
| class XRateLimitHeaderValues { | ||
| public: | ||
| const Http::LowerCaseString XRateLimitLimit{"x-ratelimit-limit"}; | ||
| const Http::LowerCaseString XRateLimitRemaining{"x-ratelimit-remaining"}; | ||
| const Http::LowerCaseString XRateLimitReset{"x-ratelimit-reset"}; | ||
|
|
||
| struct { | ||
| const std::string Window{"w"}; | ||
| const std::string Name{"name"}; | ||
| } QuotaPolicyKeys; | ||
| }; | ||
| using XRateLimitHeaders = ConstSingleton<XRateLimitHeaderValues>; | ||
|
|
||
| } // namespace RateLimit | ||
| } // namespace Common | ||
| } // namespace HttpFilters | ||
| } // namespace Extensions | ||
| } // namespace Envoy | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: maybe just
ratelimit_headers_versionorx_ratelimit_headers_version?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ok, I will changed to
x_ratelimit_headers_version.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Some thoughts, I think it is worth considering since I want to align with
ratelimit'senable_x_ratelimit_headers. Maybe we can keep the name unchanged?