0fe8b55427
Introduces the rate_limiter, a replica-side data structure meant for tracking the frequence with which each partition is being accessed (separately for reads and writes) and deciding whether the request should be accepted and processed further or rejected. The limiter is implemented as a statically allocated hashmap which keeps track of the frequency with which partitions are accessed. Its entries are incremented when an operation is admitted and are decayed exponentially over time. If a partition is detected to be accessed more than its limit allows, requests are rejected with a probability calculated in such a way that, on average, the number of accepted requests is kept at the limit. The structure currently weights a bit above 1MB and each shard is meant to keep a separate instance. All operations are O(1), including the periodic timer.
131 lines
4.3 KiB
C++
131 lines
4.3 KiB
C++
/*
|
|
* Copyright (C) 2022-present ScyllaDB
|
|
*/
|
|
|
|
/*
|
|
* SPDX-License-Identifier: AGPL-3.0-or-later
|
|
*/
|
|
|
|
#include <chrono>
|
|
#include <cstdint>
|
|
#include <seastar/core/manual_clock.hh>
|
|
#include <seastar/core/sleep.hh>
|
|
#include <seastar/core/coroutine.hh>
|
|
#include <seastar/util/later.hh>
|
|
#include <seastar/testing/test_case.hh>
|
|
|
|
#include "db/rate_limiter.hh"
|
|
|
|
using namespace seastar;
|
|
using test_rate_limiter = db::generic_rate_limiter<seastar::manual_clock>;
|
|
|
|
future<> step_seconds(int seconds) {
|
|
for (int i = 0; i < seconds; i++) {
|
|
// The rate limiter's timer executes periodically every second
|
|
// and we want the timer to run `seconds` times.
|
|
// Because `manual_clock::advance` executes each timer only once
|
|
// even if they reschedule, we cannot just advance by requested
|
|
// number of seconds - instead, we must advance multiple times
|
|
// by one second.
|
|
manual_clock::advance(std::chrono::seconds(1));
|
|
co_await yield();
|
|
}
|
|
}
|
|
|
|
SEASTAR_TEST_CASE(test_rate_limiter_no_rejections_on_sequential) {
|
|
const uint64_t token_count = 1000 * 1000;
|
|
const uint64_t limit = 1;
|
|
test_rate_limiter::label lbl;
|
|
|
|
test_rate_limiter limiter;
|
|
|
|
for (uint64_t token = 0; token < token_count; token++) {
|
|
BOOST_REQUIRE_LE(limiter.increase_and_get_counter(lbl, token), 1);
|
|
co_await maybe_yield();
|
|
}
|
|
}
|
|
|
|
SEASTAR_TEST_CASE(test_rate_limiter_partition_label_separation) {
|
|
const uint64_t token_count = 30;
|
|
const uint64_t repeat_count = 10;
|
|
std::vector<test_rate_limiter::label> labels{3};
|
|
|
|
test_rate_limiter limiter;
|
|
|
|
for (uint64_t i = 0; i < repeat_count; i++) {
|
|
for (uint64_t token = 0; token < token_count; token++) {
|
|
for (auto& l : labels) {
|
|
BOOST_REQUIRE_EQUAL(limiter.increase_and_get_counter(l, token), i + 1);
|
|
co_await maybe_yield();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
SEASTAR_TEST_CASE(test_rate_limiter_halving_over_time) {
|
|
test_rate_limiter::label lbl;
|
|
test_rate_limiter limiter;
|
|
|
|
for (int i = 0; i < 16; i++) {
|
|
limiter.increase_and_get_counter(lbl, 0);
|
|
}
|
|
|
|
// Should be cut in half
|
|
co_await step_seconds(1);
|
|
BOOST_REQUIRE_EQUAL(limiter.increase_and_get_counter(lbl, 0), (16 / 2) + 1);
|
|
|
|
// Should decrease four times (9 -> 2)
|
|
co_await step_seconds(2);
|
|
BOOST_REQUIRE_EQUAL(limiter.increase_and_get_counter(lbl, 0), (9 / 4) + 1);
|
|
|
|
// Should be reset
|
|
co_await step_seconds(10);
|
|
BOOST_REQUIRE_EQUAL(limiter.increase_and_get_counter(lbl, 0), 1);
|
|
}
|
|
|
|
SEASTAR_TEST_CASE(test_rate_limiter_time_window_wraparound_handling) {
|
|
test_rate_limiter::label lbl;
|
|
test_rate_limiter limiter;
|
|
|
|
BOOST_REQUIRE_EQUAL(limiter.increase_and_get_counter(lbl, 0), 1);
|
|
BOOST_REQUIRE_EQUAL(limiter.increase_and_get_counter(lbl, 0), 2);
|
|
BOOST_REQUIRE_EQUAL(limiter.increase_and_get_counter(lbl, 0), 3);
|
|
|
|
// Advance far into the future so that the time window wraps around
|
|
co_await step_seconds(1 << test_rate_limiter::time_window_bits);
|
|
|
|
BOOST_REQUIRE_EQUAL(limiter.increase_and_get_counter(lbl, 0), 1);
|
|
BOOST_REQUIRE_EQUAL(limiter.increase_and_get_counter(lbl, 0), 2);
|
|
BOOST_REQUIRE_EQUAL(limiter.increase_and_get_counter(lbl, 0), 3);
|
|
|
|
// TODO: Workaround for seastar#1072. Calling `manual_clock::advance`
|
|
// multiple times and then quitting the test immediately causes
|
|
// the test framework to hang. I didn't have the time to debug it, but I
|
|
// suspect there are some pending tasks which need to finish before exiting
|
|
// from the main test task.
|
|
co_await seastar::sleep(std::chrono::seconds(1));
|
|
}
|
|
|
|
SEASTAR_TEST_CASE(test_rate_limiter_account_operation) {
|
|
const uint64_t limit = 1;
|
|
const int ops_per_loop = 1000;
|
|
test_rate_limiter::label lbl;
|
|
|
|
test_rate_limiter limiter;
|
|
|
|
// We use UINT_MAX as the random parameter so that we get rejected quickly
|
|
db::per_partition_rate_limit::account_and_enforce info {
|
|
.random_variable = UINT32_MAX,
|
|
};
|
|
|
|
bool encountered_rejection = false;
|
|
for (int i = 0; i < ops_per_loop; i++) {
|
|
if (limiter.account_operation(lbl, 0, limit, info) == test_rate_limiter::can_proceed::no) {
|
|
encountered_rejection = true;
|
|
break;
|
|
}
|
|
co_await maybe_yield();
|
|
}
|
|
BOOST_REQUIRE(encountered_rejection);
|
|
}
|