889d7782ed
It's more efficient since coroutine::maybe_yield returns a lightweight struct (awaitable), not the future. Closes scylladb/scylladb#28101
131 lines
4.3 KiB
C++
131 lines
4.3 KiB
C++
/*
|
|
* Copyright (C) 2022-present ScyllaDB
|
|
*/
|
|
|
|
/*
|
|
* SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
|
|
*/
|
|
|
|
#include <chrono>
|
|
#include <cstdint>
|
|
#include <seastar/core/manual_clock.hh>
|
|
#include <seastar/core/sleep.hh>
|
|
#include <seastar/core/coroutine.hh>
|
|
#include <seastar/coroutine/maybe_yield.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;
|
|
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 coroutine::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 coroutine::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 coroutine::maybe_yield();
|
|
}
|
|
BOOST_REQUIRE(encountered_rejection);
|
|
}
|