Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: rate_limiting token bucket sub module #14

Merged
merged 7 commits into from
Apr 1, 2024

Conversation

geekbrother
Copy link
Contributor

@geekbrother geekbrother commented Mar 25, 2024

Description

The rate-limiting token bucket mechanism (based on #231) uses Redis and Moka in-memory cache.

The initial implementation was updated to use the local app cache in Moka and omit the Redis RTT calls (for single key calls) because we know the TTL when the token will be refilled in cases when the key is rate-limited. So we don't need to call the Redis until the TTL expires and we can use the local memory Moka cache. This should significantly eliminate Redis RTT calls in the case of DOS/flood. The initial Lua script was a bit optimized and minified.

How Has This Been Tested?

Due Diligence

  • Breaking change
  • Requires a documentation update
  • Requires a e2e/integration test update

@geekbrother geekbrother self-assigned this Mar 25, 2024
@geekbrother geekbrother force-pushed the feat/rate_limiting_lib branch 2 times, most recently from 21506a7 to 4679aba Compare March 25, 2024 11:59
@geekbrother geekbrother force-pushed the feat/rate_limiting_lib branch from 4679aba to b0889c1 Compare March 25, 2024 12:05
@geekbrother geekbrother requested a review from chris13524 March 25, 2024 22:10
@geekbrother geekbrother changed the title feat: rate_limiting library feat: rate_limiting token bucket sub module Mar 25, 2024
@geekbrother geekbrother marked this pull request as ready for review March 25, 2024 22:15
}

#[tokio::test]
async fn test_token_bucket_many() {
Copy link
Member

@chris13524 chris13524 Mar 26, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test has significantly less coverage than the one from Notify Server. Can we copy those cases directly so we can keep the level of coverage?

Missing:

  • refill_in calculation
  • checking that tokens recover 1 at a time rather than all at once
  • test for separate keys

Why did you remove the ability to mock the clock? That helps ensure that tests aren't flaky as well as I think making the refill_in test case possible.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense to update the test coverage with the missing ones, thanks for pointing out this!

Why did you remove the ability to mock the clock? That helps ensure that tests aren't flaky as well as I think making the refill_in test case possible.

Maybe it's better to pass the timestamp to the function instead of hacking the clock for the test (it looks a bit hacky to me). We can just pass Utc::now().timestamp_millis() to the function instead of getting it inside of it.

I have updated tests and the function in the 729ec96.

The test coverage now has:

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are you planning to replace the Notify Server implementation with this standard one? If so, the mocked clock is was necessary to be able to test that rate limits are applied to regular endpoints.

It's common practice to mock clocks, I don't see what's "hacky" about that.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I want to apply this unified implementation to the Blockchain-api and Echo server and then probably to the Notify (not on the table yet).

@geekbrother geekbrother marked this pull request as draft March 27, 2024 13:23
@geekbrother geekbrother marked this pull request as ready for review March 27, 2024 15:46
@geekbrother geekbrother requested a review from chris13524 March 27, 2024 15:46
@geekbrother geekbrother marked this pull request as draft March 27, 2024 17:51
@geekbrother geekbrother marked this pull request as ready for review March 27, 2024 21:42
@@ -6,39 +6,24 @@ local refillRate = tonumber(ARGV[3]) -- how many tokens are refilled after each
local now = tonumber(ARGV[4]) -- current timestamp in milliseconds

local results = {}

for i, key in ipairs(keys) do
for i, key in ipairs(KEYS) do
Copy link
Member

@chris13524 chris13524 Mar 27, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

☹️ I would prefer if we touched the Lua script as minimally as possible. It's a direct copy of the linked script above with only the tweaks necessary to make it work with multiple keys. You may have introduced bugs and made it harder to keep in-sync with the implementation above. I found a bug before in the implementation and contributed it back. Do you really think you've made significant performance improvements worth risking these changes? Looks to me that these changes might actually be less performant due to the removal of locals and I don't think one-time complication or hashing speed really has an effect.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reverted it in b9ecde9 if you find it hard to debug the minified version of the Lua script since it's not a huge optimization. Let's stick to the initial one and move forward with it.

crates/rate_limit/src/token_bucket.lua Outdated Show resolved Hide resolved
crates/rate_limit/src/token_bucket.lua Outdated Show resolved Hide resolved
crates/rate_limit/src/token_bucket.lua Outdated Show resolved Hide resolved
@geekbrother geekbrother requested a review from nopestack March 28, 2024 09:20
@geekbrother geekbrother force-pushed the feat/rate_limiting_lib branch from cf32287 to b9ecde9 Compare March 28, 2024 10:00
@geekbrother geekbrother requested a review from chris13524 March 28, 2024 10:03
/// Rate limit check using a token bucket algorithm for one key and in-memory
/// cache for rate-limited keys. `mem_cache` TTL must be set to the same value
/// as the refill interval.
pub async fn token_bucket(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

mem_cache TTL must be set to the same value as the refill interval.

What does this mean? Does Cache support different keys with different refill intervals e.g. for what Notify Server needs?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The TTL is for all records. If there are requirements to refill different rates for different cases it can be produced with the single refill interval (lowest) and different refill rates by tuning the refill rate to the single refill interval, without adding additional TTL per key.

@geekbrother geekbrother merged commit e2bc55d into main Apr 1, 2024
7 checks passed
@geekbrother geekbrother deleted the feat/rate_limiting_lib branch April 1, 2024 09:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants