This is an automated email from the ASF dual-hosted git repository.
bobhan1 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new c81ad5b8406 [improvement](cloud) Add dry-run mode for BE-to-MS RPC
rate limiting (#66977)
c81ad5b8406 is described below
commit c81ad5b84063fc78158a8426cce8962295a84c35
Author: bobhan1 <[email protected]>
AuthorDate: Fri Aug 28 10:14:12 2026 +0800
[improvement](cloud) Add dry-run mode for BE-to-MS RPC rate limiting
(#66977)
### What problem does this PR solve?
Issue Number: N/A
Related PR: #66969, #66940
Problem Summary: BE-to-MetaService host-level and table-level rate
limiting needs production-safe observability before thresholds are
enforced. This PR adds dry-run switches for both paths. Dry-run keeps
limit evaluation, would-wait metrics, per-table QPS collection, and
table-level MS_BUSY state transitions active, but does not delay
requests.
| Scope | New config | Default | Existing enforcement config | Purpose |
| --- | --- | --- | --- | --- |
| Host-level | `enable_ms_rpc_host_level_rate_limit_dry_run` | `true` |
`enable_ms_rpc_host_level_rate_limit` (default `false`) | Evaluate each
configured per-RPC token bucket, record would-wait observations, and log
would-throttle decisions without sleeping. |
| Table-level | `enable_ms_backpressure_handling_dry_run` | `true` |
`enable_ms_backpressure_handling` (default `false`) | Keep per-table QPS
collection and the MS_BUSY adaptive-throttling state machine active, but
do not sleep on a throttle decision. |
For each scope, the dry-run and enforcement configs interact as follows:
- When both configs are `false`, rate-limit evaluation for that scope is
disabled.
- When only the enforcement config is `true`, the limiter is enforced
and the request sleeps for the calculated wait time.
- When the dry-run config is `true`, dry-run takes precedence regardless
of the enforcement config: the limiter state, metrics, and logs are
updated, but the request does not sleep.
Both host-level and table-level dry-run paths reserve against the same
limiter state used by enforcement, matching the FE dry-run model:
dry-run suppresses sleeping but preserves reservation state. When a
table-level downgrade restores a previous QPS limit, it resets queued
reservations so stale wait debt does not carry into the relaxed limit;
the final downgrade removes the limiter.
The table-level coordinator uses 64-bit counters with overflow-free
saturation at each timing threshold. Its MS_BUSY downgrade timer is
reset to inactive when there is no pending upgrade history, while the
upgrade cooldown counter is preserved and stops at the configured
threshold so runtime cooldown changes retain the existing timing
semantics.
### Signals after rate limiting is triggered
Throttle-decision logs use `INFO` level and are suppressed independently
per RPC type to at most one log per second. They are emitted only when
the limiter calculates a positive wait.
Host-level dry-run logs explicitly report the estimated wait:
```text
[ms-throttle] host-level rate limiter dry run would throttle MS RPC
request, rpc=get rowset, estimated_wait_ns=1250000, qps_limit=320
```
With enforcement enabled and dry-run disabled, the corresponding log
reports the actual sleep:
```text
[ms-throttle] host-level rate limiter throttled MS RPC request, rpc=get
rowset, sleep_ns=1250000, qps_limit=320
```
Table-level dry-run and enforcement use the same trigger log so
operators can search one stable pattern:
```text
[ms-throttle] table-level rate limiter triggered for MS RPC request,
rpc=commit_rowset, table_id=10001, wait_us=8200, current_qps=146.7,
qps_limit=100
```
For table-level dry-run, `wait_us` is the estimated wait and the request
is not delayed. With enforcement enabled and dry-run disabled, the
request sleeps for that wait. MS_BUSY-driven table-level state
transitions also emit `INFO` logs, for example:
```text
[ms-throttle] received MS_BUSY, triggering upgrade
[ms-throttle] upgrade: rpc=commit_rowset, table_id=10001, current_qps=200,
old_limit=0, new_limit=150
[ms-throttle] downgrade: rpc=commit_rowset, table_id=10001, removed limit
```
The same events are observable through bvars: host-level per-RPC latency
recorders based on `host_level_ms_rpc_rate_limit_sleep`, table-level
`ms_rpc_backpressure_throttle_wait_<rpc>` latency recorders, and
`ms_rpc_backpressure_ms_busy_*`, `ms_rpc_backpressure_upgrade_*`, and
`ms_rpc_backpressure_downgrade_*` counters/windows.
### Release note
Add dry-run observation for BE-to-MetaService host-level and table-level
RPC rate limiting without delaying requests. Both dry-run switches are
enabled by default and take precedence over enforcement. Dry-run and
enforcement share limiter reservation state, table-level downgrade
resets queued reservations when restoring a previous QPS limit, and
long-running table-level coordinator timers no longer overflow.
### Check List (For Author)
- Test
- [ ] Regression test
- [x] Unit Test
- `./run-be-ut.sh --run
--filter='HostLevelMSRpcRateLimitersTest.*:StrictQpsLimiterTest.*:TableRpcThrottlerTest.*:MSBackpressureHandlerTest.*:RpcThrottleStateMachineTest.*:RpcThrottleIntegrationTest.*'
-j100` (56 tests passed)
- `./run-be-ut.sh --run
--filter='TokenBucketRateLimiterTest.*:S3RateLimiterManagerTest.*:S3RateLimiterMetricsTest.*'
-j100` (11 tests passed)
- `./run-be-ut.sh --run
--filter='MSBackpressureHandlerTest.*:RpcThrottleStateMachineTest.*:RpcThrottleCoordinatorTest.*:RpcThrottleIntegrationTest.*'
-j100` (44 tests passed)
- `./build.sh --be -j100`
- `build-support/check-format.sh`
- `build-support/check-build-hygiene.sh`
- [ ] Manual test (add detailed scripts or steps below)
- [ ] No need to test or manual test. Explain why:
- [ ] This is a refactor/code format and no logic has been changed.
- [ ] Previous test can cover this change.
- [ ] No code files have been changed.
- [ ] Other reason
- Behavior changed:
- [ ] No.
- [x] Yes. Dry-run observation for both limiter levels is enabled by
default and does not delay MetaService RPCs; dry-run and enforcement
share reservation state; table-level downgrade resets queued
reservations; coordinator timing counters saturate without overflowing;
actual and dry-run decisions emit per-RPC rate-limited `INFO` logs.
- Does this need documentation?
- [x] No.
- [ ] Yes.
### Check List (For Reviewer who merge this PR)
- [ ] Confirm the release note
- [ ] Confirm test cases
- [ ] Confirm document
- [ ] Add branch pick label
---
be/src/cloud/cloud_meta_mgr.cpp | 70 +++++---
be/src/cloud/cloud_ms_backpressure_handler.cpp | 197 ++++++++++++++++-----
be/src/cloud/cloud_ms_backpressure_handler.h | 56 ++++--
be/src/cloud/cloud_ms_rpc_rate_limiters.cpp | 28 ++-
be/src/cloud/cloud_ms_rpc_rate_limiters.h | 10 +-
be/src/cloud/cloud_throttle_state_machine.cpp | 74 +++++---
be/src/cloud/cloud_throttle_state_machine.h | 16 +-
be/src/cloud/config.cpp | 2 +
be/src/cloud/config.h | 8 +
.../cloud/cloud_ms_backpressure_handler_test.cpp | 168 +++++++++++++++++-
be/test/cloud/cloud_ms_rpc_rate_limiters_test.cpp | 63 ++++++-
.../cloud/cloud_throttle_state_machine_test.cpp | 89 +++++++++-
common/cpp/token_bucket_rate_limiter.cpp | 27 ++-
common/cpp/token_bucket_rate_limiter.h | 8 +
.../cloud/test_cloud_ms_rpc_table_throttle.groovy | 1 +
15 files changed, 685 insertions(+), 132 deletions(-)
diff --git a/be/src/cloud/cloud_meta_mgr.cpp b/be/src/cloud/cloud_meta_mgr.cpp
index a13b72fa84c..6533734d00c 100644
--- a/be/src/cloud/cloud_meta_mgr.cpp
+++ b/be/src/cloud/cloud_meta_mgr.cpp
@@ -73,6 +73,7 @@
#include "util/network_util.h"
#include "util/s3_util.h"
#include "util/thrift_rpc_helper.h"
+#include "util/time.h"
namespace doris::cloud {
using namespace ErrorCode;
@@ -494,33 +495,54 @@ struct RpcRateLimitCtx {
int64_t table_id {-1}; // For table-level backpressure, passed from caller
};
-// Apply rate limiting before RPC (both host-level and table-level)
-void apply_rate_limit(MetaServiceRPC rpc, const RpcRateLimitCtx& ctx) {
- // Table-level rate limit (for load-related RPCs only)
- if (ctx.backpressure_handler && ctx.table_id > 0) {
- LoadRelatedRpc load_rpc = to_load_related_rpc(rpc);
- if (load_rpc != LoadRelatedRpc::COUNT) {
- auto wait_until = ctx.backpressure_handler->before_rpc(load_rpc,
ctx.table_id);
- auto now = std::chrono::steady_clock::now();
- if (wait_until > now) {
- auto wait_us =
-
std::chrono::duration_cast<std::chrono::microseconds>(wait_until - now)
- .count();
- if (wait_us > 0) {
- if (auto* recorder = get_throttle_wait_recorder(load_rpc);
- recorder != nullptr) {
- *recorder << wait_us;
- }
- bthread_usleep(wait_us);
- }
- }
- }
+void apply_table_level_rate_limit(MetaServiceRPC rpc, const RpcRateLimitCtx&
ctx) {
+ if (ctx.backpressure_handler == nullptr || ctx.table_id <= 0) {
+ return;
}
- // Host-level rate limit
- if (ctx.host_limiters) {
- ctx.host_limiters->limit(rpc);
+ const auto load_rpc = to_load_related_rpc(rpc);
+ if (load_rpc == LoadRelatedRpc::COUNT) {
+ return;
+ }
+
+ const auto decision = ctx.backpressure_handler->before_rpc(load_rpc,
ctx.table_id);
+ const auto now = std::chrono::steady_clock::now();
+ if (decision.wait_until <= now) {
+ return;
+ }
+
+ const auto wait_us =
+
std::chrono::duration_cast<std::chrono::microseconds>(decision.wait_until - now)
+ .count();
+ if (wait_us <= 0) {
+ return;
+ }
+
+ auto* recorder = get_throttle_wait_recorder(load_rpc);
+ DCHECK(recorder);
+ *recorder << wait_us;
+ if (ctx.backpressure_handler->should_log_throttle(load_rpc,
MonotonicMicros())) {
+ const double current_qps =
+ ctx.backpressure_handler->get_current_qps(load_rpc,
ctx.table_id);
+ LOG(INFO) << "[ms-throttle] table-level rate limiter triggered for MS
RPC request"
+ << ", rpc=" << load_related_rpc_name(load_rpc) << ",
table_id=" << ctx.table_id
+ << ", wait_us=" << wait_us << ", current_qps=" << current_qps
+ << ", qps_limit=" << decision.qps_limit;
+ }
+
+ if (decision.dry_run) {
+ return;
+ }
+ bthread_usleep(wait_us);
+}
+
+// Apply rate limiting before RPC (both host-level and table-level)
+void apply_rate_limit(MetaServiceRPC rpc, const RpcRateLimitCtx& ctx) {
+ apply_table_level_rate_limit(rpc, ctx);
+ if (ctx.host_limiters == nullptr) {
+ return;
}
+ ctx.host_limiters->limit(rpc);
}
// Record RPC QPS statistics after RPC (for table-level tracking)
diff --git a/be/src/cloud/cloud_ms_backpressure_handler.cpp
b/be/src/cloud/cloud_ms_backpressure_handler.cpp
index e99cc281001..7ccc3853a71 100644
--- a/be/src/cloud/cloud_ms_backpressure_handler.cpp
+++ b/be/src/cloud/cloud_ms_backpressure_handler.cpp
@@ -28,9 +28,23 @@
#include "cloud/config.h"
#include "common/status.h"
#include "util/thread.h"
+#include "util/time.h"
namespace doris::cloud {
+namespace {
+
+constexpr std::chrono::milliseconds kQpsRegistryCleanupInterval =
std::chrono::minutes(1);
+constexpr std::chrono::milliseconds kQpsRegistryMinInactiveTimeout =
std::chrono::minutes(10);
+
+std::chrono::milliseconds qps_registry_inactive_timeout() {
+ return std::max(kQpsRegistryMinInactiveTimeout,
+ std::chrono::duration_cast<std::chrono::milliseconds>(
+
std::chrono::seconds(config::ms_rpc_table_qps_window_sec)));
+}
+
+} // namespace
+
// Global bvar metrics
bvar::Adder<uint64_t>
g_backpressure_upgrade_count("ms_rpc_backpressure_upgrade_count");
bvar::Window<bvar::Adder<uint64_t>>
g_backpressure_upgrade_60s("ms_rpc_backpressure_upgrade_60s",
@@ -88,12 +102,15 @@ StrictQpsLimiter::Clock::time_point
StrictQpsLimiter::reserve() {
return result;
}
-void StrictQpsLimiter::update_qps(double new_qps) {
+void StrictQpsLimiter::update_qps(double new_qps, bool reset_reservation) {
if (new_qps <= 0) {
new_qps = 1.0;
}
std::lock_guard lock(_mtx);
_interval_ns = static_cast<int64_t>(1e9 / new_qps);
+ if (reset_reservation) {
+ _next_allowed_time = Clock::now();
+ }
}
double StrictQpsLimiter::get_qps() const {
@@ -115,6 +132,7 @@ TableRpcQpsCounter::TableRpcQpsCounter(int64_t table_id,
LoadRelatedRpc rpc_type
}
void TableRpcQpsCounter::increment() {
+ _last_record_time_us.store(MonotonicMicros(), std::memory_order_relaxed);
(*_counter) << 1;
}
@@ -124,27 +142,44 @@ double TableRpcQpsCounter::get_qps() const {
// ============== TableRpcQpsRegistry ==============
-TableRpcQpsRegistry::TableRpcQpsRegistry() = default;
+TableRpcQpsRegistry::TableRpcQpsRegistry()
+ : TableRpcQpsRegistry(kQpsRegistryCleanupInterval,
qps_registry_inactive_timeout()) {}
-void TableRpcQpsRegistry::record(LoadRelatedRpc rpc_type, int64_t table_id) {
- auto* counter = get_or_create_counter(rpc_type, table_id);
- if (counter) {
- counter->increment();
+TableRpcQpsRegistry::TableRpcQpsRegistry(std::chrono::milliseconds
cleanup_interval,
+ std::chrono::milliseconds
inactive_timeout)
+ : _cleanup_interval(cleanup_interval),
+ _inactive_timeout(inactive_timeout),
+ _cleanup_stop_latch(1) {
+ DORIS_CHECK_GT(_cleanup_interval.count(), 0);
+ DORIS_CHECK_GE(_inactive_timeout.count(), 0);
+
+ auto st = Thread::create(
+ "TableRpcQpsRegistry", "cleanup_thread", [this]() {
this->_cleanup_thread_callback(); },
+ &_cleanup_thread);
+ if (!st.ok()) {
+ LOG(WARNING) << "[ms-throttle] failed to create table QPS registry
cleanup thread: " << st;
}
}
-TableRpcQpsCounter* TableRpcQpsRegistry::get_or_create_counter(LoadRelatedRpc
rpc_type,
- int64_t
table_id) {
+TableRpcQpsRegistry::~TableRpcQpsRegistry() {
+ _cleanup_stop_latch.count_down();
+ if (_cleanup_thread) {
+ _cleanup_thread->join();
+ }
+}
+
+void TableRpcQpsRegistry::record(LoadRelatedRpc rpc_type, int64_t table_id) {
size_t idx = static_cast<size_t>(rpc_type);
if (idx >= static_cast<size_t>(LoadRelatedRpc::COUNT)) {
- return nullptr;
+ return;
}
{
std::shared_lock lock(_mutex);
auto it = _counters[idx].find(table_id);
if (it != _counters[idx].end()) {
- return it->second.get();
+ it->second->increment();
+ return;
}
}
@@ -152,14 +187,14 @@ TableRpcQpsCounter*
TableRpcQpsRegistry::get_or_create_counter(LoadRelatedRpc rp
// Double check after acquiring exclusive lock
auto it = _counters[idx].find(table_id);
if (it != _counters[idx].end()) {
- return it->second.get();
+ it->second->increment();
+ return;
}
auto counter = std::make_unique<TableRpcQpsCounter>(table_id, rpc_type,
config::ms_rpc_table_qps_window_sec);
- auto* ptr = counter.get();
+ counter->increment();
_counters[idx][table_id] = std::move(counter);
- return ptr;
}
std::vector<std::pair<int64_t, double>> TableRpcQpsRegistry::get_top_k_tables(
@@ -216,25 +251,74 @@ double TableRpcQpsRegistry::get_qps(LoadRelatedRpc
rpc_type, int64_t table_id) c
return 0;
}
-void TableRpcQpsRegistry::cleanup_inactive_tables() {
- std::unique_lock lock(_mutex);
+size_t TableRpcQpsRegistry::cleanup_inactive_tables() {
+ const int64_t inactive_before_us =
+ MonotonicMicros() -
+
std::chrono::duration_cast<std::chrono::microseconds>(_inactive_timeout).count();
+ std::array<std::vector<int64_t>,
static_cast<size_t>(LoadRelatedRpc::COUNT)> candidates;
- for (size_t idx = 0; idx < static_cast<size_t>(LoadRelatedRpc::COUNT);
++idx) {
- auto& counter_map = _counters[idx];
- for (auto it = counter_map.begin(); it != counter_map.end();) {
- // Remove counters with zero QPS for a long time
- if (it->second->get_qps() < 0.01) {
- it = counter_map.erase(it);
- } else {
- ++it;
+ {
+ std::shared_lock lock(_mutex);
+ for (size_t idx = 0; idx < static_cast<size_t>(LoadRelatedRpc::COUNT);
++idx) {
+ for (const auto& [table_id, counter] : _counters[idx]) {
+ if (counter->last_record_time_us() <= inactive_before_us) {
+ candidates[idx].push_back(table_id);
+ }
}
}
}
+
+ size_t candidate_count = 0;
+ for (const auto& rpc_candidates : candidates) {
+ candidate_count += rpc_candidates.size();
+ }
+
+ std::vector<std::unique_ptr<TableRpcQpsCounter>> counters_to_destroy;
+ counters_to_destroy.reserve(candidate_count);
+ {
+ std::unique_lock lock(_mutex);
+ for (size_t idx = 0; idx < static_cast<size_t>(LoadRelatedRpc::COUNT);
++idx) {
+ auto& counter_map = _counters[idx];
+ for (int64_t table_id : candidates[idx]) {
+ auto it = counter_map.find(table_id);
+ if (it != counter_map.end() &&
+ it->second->last_record_time_us() <= inactive_before_us) {
+ counters_to_destroy.push_back(std::move(it->second));
+ counter_map.erase(it);
+ }
+ }
+ }
+ }
+ return counters_to_destroy.size();
+}
+
+size_t TableRpcQpsRegistry::get_tracked_table_count(LoadRelatedRpc rpc_type)
const {
+ size_t idx = static_cast<size_t>(rpc_type);
+ if (idx >= static_cast<size_t>(LoadRelatedRpc::COUNT)) {
+ return 0;
+ }
+
+ std::shared_lock lock(_mutex);
+ return _counters[idx].size();
+}
+
+void TableRpcQpsRegistry::_cleanup_thread_callback() {
+ while (!_cleanup_stop_latch.wait_for(_cleanup_interval)) {
+ size_t removed = cleanup_inactive_tables();
+ if (removed > 0) {
+ LOG(INFO) << "[ms-throttle] cleaned up inactive table QPS
counters: removed="
+ << removed;
+ }
+ }
}
// ============== TableRpcThrottler ==============
TableRpcThrottler::TableRpcThrottler() {
+ for (auto& next_log_time_us : _next_log_time_us) {
+ next_log_time_us.store(0, std::memory_order_relaxed);
+ }
+
// Initialize bvar for throttled table counts
for (size_t i = 0; i < static_cast<size_t>(LoadRelatedRpc::COUNT); ++i) {
std::string bvar_name =
fmt::format("ms_rpc_backpressure_throttled_tables_{}",
@@ -245,15 +329,35 @@ TableRpcThrottler::TableRpcThrottler() {
std::chrono::steady_clock::time_point
TableRpcThrottler::throttle(LoadRelatedRpc rpc_type,
int64_t
table_id) {
+ return throttle(rpc_type, table_id, false).wait_until;
+}
+
+TableRpcThrottleDecision TableRpcThrottler::throttle(LoadRelatedRpc rpc_type,
int64_t table_id,
+ bool dry_run) {
std::shared_lock lock(_mutex);
auto it = _limiters.find({rpc_type, table_id});
if (it == _limiters.end()) {
- return std::chrono::steady_clock::now();
+ return {.wait_until = std::chrono::steady_clock::now(), .dry_run =
dry_run};
}
- return it->second->reserve();
+ return {
+ .wait_until = it->second->reserve(),
+ .qps_limit = it->second->get_qps(),
+ .dry_run = dry_run,
+ };
}
-void TableRpcThrottler::set_qps_limit(LoadRelatedRpc rpc_type, int64_t
table_id, double qps_limit) {
+bool TableRpcThrottler::should_log(LoadRelatedRpc rpc_type, int64_t now_us) {
+ size_t idx = static_cast<size_t>(rpc_type);
+ DCHECK_LT(idx, static_cast<size_t>(LoadRelatedRpc::COUNT));
+ auto& next_log_time_us = _next_log_time_us[idx];
+ int64_t expected = next_log_time_us.load(std::memory_order_relaxed);
+ return now_us >= expected &&
+ next_log_time_us.compare_exchange_strong(expected, now_us +
MICROS_PER_SEC,
+ std::memory_order_relaxed);
+}
+
+void TableRpcThrottler::set_qps_limit(LoadRelatedRpc rpc_type, int64_t
table_id, double qps_limit,
+ bool reset_reservation) {
if (qps_limit <= 0) {
return;
}
@@ -262,7 +366,7 @@ void TableRpcThrottler::set_qps_limit(LoadRelatedRpc
rpc_type, int64_t table_id,
auto key = std::make_pair(rpc_type, table_id);
auto it = _limiters.find(key);
if (it != _limiters.end()) {
- it->second->update_qps(qps_limit);
+ it->second->update_qps(qps_limit, reset_reservation);
} else {
_limiters[key] = std::make_unique<StrictQpsLimiter>(qps_limit);
// Update bvar count
@@ -384,14 +488,15 @@ MSBackpressureHandler::~MSBackpressureHandler() {
void MSBackpressureHandler::_tick_thread_callback() {
// Fixed tick interval: 1 second. Since 1 tick = 1 ms, advance by 1000
ticks each iteration.
- constexpr int kTickIntervalMs = 1000;
+ constexpr int64_t kTickIntervalMs = 1000;
while (!_stop_latch.wait_for(std::chrono::milliseconds(kTickIntervalMs))) {
_advance_time(kTickIntervalMs);
}
}
-void MSBackpressureHandler::_advance_time(int ticks) {
- if (!config::enable_ms_backpressure_handling) {
+void MSBackpressureHandler::_advance_time(int64_t ticks) {
+ if (!config::enable_ms_backpressure_handling &&
+ !config::enable_ms_backpressure_handling_dry_run) {
return;
}
@@ -413,7 +518,8 @@ void MSBackpressureHandler::_advance_time(int ticks) {
bool MSBackpressureHandler::on_ms_busy() {
g_ms_busy_count << 1;
- if (!config::enable_ms_backpressure_handling) {
+ if (!config::enable_ms_backpressure_handling &&
+ !config::enable_ms_backpressure_handling_dry_run) {
return false;
}
@@ -440,17 +546,27 @@ bool MSBackpressureHandler::on_ms_busy() {
return true;
}
-std::chrono::steady_clock::time_point
MSBackpressureHandler::before_rpc(LoadRelatedRpc rpc_type,
-
int64_t table_id) {
- if (!config::enable_ms_backpressure_handling) {
- return std::chrono::steady_clock::now();
+TableRpcThrottleDecision MSBackpressureHandler::before_rpc(LoadRelatedRpc
rpc_type,
+ int64_t table_id) {
+ const bool dry_run = config::enable_ms_backpressure_handling_dry_run;
+ if (!config::enable_ms_backpressure_handling && !dry_run) {
+ return {.wait_until = std::chrono::steady_clock::now()};
}
- return _throttler->throttle(rpc_type, table_id);
+ return _throttler->throttle(rpc_type, table_id, dry_run);
+}
+
+bool MSBackpressureHandler::should_log_throttle(LoadRelatedRpc rpc_type,
int64_t now_us) {
+ return _throttler->should_log(rpc_type, now_us);
+}
+
+double MSBackpressureHandler::get_current_qps(LoadRelatedRpc rpc_type, int64_t
table_id) const {
+ return _qps_registry->get_qps(rpc_type, table_id);
}
void MSBackpressureHandler::after_rpc(LoadRelatedRpc rpc_type, int64_t
table_id) {
- if (!config::enable_ms_backpressure_handling) {
+ if (!config::enable_ms_backpressure_handling &&
+ !config::enable_ms_backpressure_handling_dry_run) {
return;
}
@@ -479,11 +595,11 @@ size_t MSBackpressureHandler::upgrade_level() const {
return _state_machine->upgrade_level();
}
-int MSBackpressureHandler::ticks_since_last_ms_busy() const {
+int64_t MSBackpressureHandler::ticks_since_last_ms_busy() const {
return _coordinator->ticks_since_last_ms_busy();
}
-int MSBackpressureHandler::ticks_since_last_upgrade() const {
+int64_t MSBackpressureHandler::ticks_since_last_upgrade() const {
return _coordinator->ticks_since_last_upgrade();
}
@@ -491,7 +607,8 @@ void MSBackpressureHandler::_apply_actions(const
std::vector<RpcThrottleAction>&
for (const auto& action : actions) {
switch (action.type) {
case RpcThrottleAction::Type::SET_LIMIT:
- _throttler->set_qps_limit(action.rpc_type, action.table_id,
action.qps_limit);
+ _throttler->set_qps_limit(action.rpc_type, action.table_id,
action.qps_limit,
+ action.reset_reservation);
break;
case RpcThrottleAction::Type::REMOVE_LIMIT:
_throttler->remove_qps_limit(action.rpc_type, action.table_id);
diff --git a/be/src/cloud/cloud_ms_backpressure_handler.h
b/be/src/cloud/cloud_ms_backpressure_handler.h
index f32100c2c46..8d5025cf641 100644
--- a/be/src/cloud/cloud_ms_backpressure_handler.h
+++ b/be/src/cloud/cloud_ms_backpressure_handler.h
@@ -20,6 +20,7 @@
#include <bvar/bvar.h>
#include <array>
+#include <atomic>
#include <chrono>
#include <map>
#include <memory>
@@ -48,8 +49,8 @@ public:
// Caller should sleep until this time point
Clock::time_point reserve();
- // Dynamically update the QPS limit
- void update_qps(double new_qps);
+ // Dynamically update the QPS limit, optionally discarding queued
reservations.
+ void update_qps(double new_qps, bool reset_reservation = false);
// Get current QPS limit
double get_qps() const;
@@ -72,6 +73,10 @@ public:
// Get current QPS (average over the configured time window)
double get_qps() const;
+ int64_t last_record_time_us() const {
+ return _last_record_time_us.load(std::memory_order_relaxed);
+ }
+
int64_t table_id() const { return _table_id; }
LoadRelatedRpc rpc_type() const { return _rpc_type; }
@@ -81,13 +86,16 @@ private:
std::unique_ptr<bvar::Adder<int64_t>> _counter;
std::unique_ptr<bvar::PerSecond<bvar::Adder<int64_t>>> _qps;
+ std::atomic<int64_t> _last_record_time_us {0};
};
// Registry managing QPS counters for all tables
class TableRpcQpsRegistry {
public:
TableRpcQpsRegistry();
- ~TableRpcQpsRegistry() = default;
+ TableRpcQpsRegistry(std::chrono::milliseconds cleanup_interval,
+ std::chrono::milliseconds inactive_timeout);
+ ~TableRpcQpsRegistry();
// Record one RPC call for the given table
void record(LoadRelatedRpc rpc_type, int64_t table_id);
@@ -100,11 +108,12 @@ public:
double get_qps(LoadRelatedRpc rpc_type, int64_t table_id) const;
// Clean up counters for tables that have been inactive for a long time
- void cleanup_inactive_tables();
+ size_t cleanup_inactive_tables();
+
+ size_t get_tracked_table_count(LoadRelatedRpc rpc_type) const;
private:
- // Get or create counter for (rpc_type, table_id)
- TableRpcQpsCounter* get_or_create_counter(LoadRelatedRpc rpc_type, int64_t
table_id);
+ void _cleanup_thread_callback();
mutable std::shared_mutex _mutex;
@@ -112,6 +121,17 @@ private:
std::array<std::unordered_map<int64_t,
std::unique_ptr<TableRpcQpsCounter>>,
static_cast<size_t>(LoadRelatedRpc::COUNT)>
_counters;
+
+ const std::chrono::milliseconds _cleanup_interval;
+ const std::chrono::milliseconds _inactive_timeout;
+ std::shared_ptr<Thread> _cleanup_thread;
+ CountDownLatch _cleanup_stop_latch;
+};
+
+struct TableRpcThrottleDecision {
+ std::chrono::steady_clock::time_point wait_until;
+ double qps_limit {0};
+ bool dry_run {false};
};
// Table-level throttler managing StrictQpsLimiter for each (RPC type, table)
pair
@@ -123,9 +143,14 @@ public:
// Called before RPC execution, returns the time point when execution is
allowed
// Returns now if no limit is set
std::chrono::steady_clock::time_point throttle(LoadRelatedRpc rpc_type,
int64_t table_id);
+ TableRpcThrottleDecision throttle(LoadRelatedRpc rpc_type, int64_t
table_id, bool dry_run);
+
+ // Log suppression is independent for every RPC type.
+ bool should_log(LoadRelatedRpc rpc_type, int64_t now_us);
// Set or update the QPS limit for a table
- void set_qps_limit(LoadRelatedRpc rpc_type, int64_t table_id, double
qps_limit);
+ void set_qps_limit(LoadRelatedRpc rpc_type, int64_t table_id, double
qps_limit,
+ bool reset_reservation = false);
// Remove the QPS limit for a table
void remove_qps_limit(LoadRelatedRpc rpc_type, int64_t table_id);
@@ -149,9 +174,10 @@ public:
private:
mutable std::shared_mutex _mutex;
- // (rpc_type, table_id) -> StrictQpsLimiter
std::map<std::pair<LoadRelatedRpc, int64_t>,
std::unique_ptr<StrictQpsLimiter>> _limiters;
+ std::array<std::atomic<int64_t>,
static_cast<size_t>(LoadRelatedRpc::COUNT)> _next_log_time_us;
+
// bvar: current throttled table count per RPC type
std::array<std::unique_ptr<bvar::Status<size_t>>,
static_cast<size_t>(LoadRelatedRpc::COUNT)>
_throttled_table_counts;
@@ -168,9 +194,11 @@ public:
// Returns true if throttle upgrade was triggered
bool on_ms_busy();
- // Called before RPC execution, performs throttle wait
- // Returns the time point to wait until
- std::chrono::steady_clock::time_point before_rpc(LoadRelatedRpc rpc_type,
int64_t table_id);
+ // Called before RPC execution and returns the actual or dry-run throttle
decision.
+ TableRpcThrottleDecision before_rpc(LoadRelatedRpc rpc_type, int64_t
table_id);
+
+ bool should_log_throttle(LoadRelatedRpc rpc_type, int64_t now_us);
+ double get_current_qps(LoadRelatedRpc rpc_type, int64_t table_id) const;
// Called after RPC execution, records QPS statistics
void after_rpc(LoadRelatedRpc rpc_type, int64_t table_id);
@@ -184,15 +212,15 @@ public:
// Query current state
size_t upgrade_level() const;
- int ticks_since_last_ms_busy() const;
- int ticks_since_last_upgrade() const;
+ int64_t ticks_since_last_ms_busy() const;
+ int64_t ticks_since_last_upgrade() const;
private:
// Background thread that periodically advances time
void _tick_thread_callback();
// Advance time by specified ticks, handle any triggered events (e.g.,
downgrade)
- void _advance_time(int ticks);
+ void _advance_time(int64_t ticks);
// Apply actions to the throttler
void _apply_actions(const std::vector<RpcThrottleAction>& actions);
diff --git a/be/src/cloud/cloud_ms_rpc_rate_limiters.cpp
b/be/src/cloud/cloud_ms_rpc_rate_limiters.cpp
index 74fe85c575e..dd741b47ad4 100644
--- a/be/src/cloud/cloud_ms_rpc_rate_limiters.cpp
+++ b/be/src/cloud/cloud_ms_rpc_rate_limiters.cpp
@@ -23,6 +23,7 @@
#include "cloud/config.h"
#include "util/cpu_info.h"
+#include "util/time.h"
namespace doris::cloud {
@@ -86,6 +87,13 @@ RpcRateLimiter::RpcRateLimiter(int qps, std::string_view
op_name) {
});
}
+bool RpcRateLimiter::should_log(int64_t now_us) {
+ int64_t next_log_time_us =
_next_log_time_us.load(std::memory_order_relaxed);
+ return now_us >= next_log_time_us &&
+ _next_log_time_us.compare_exchange_strong(next_log_time_us, now_us
+ MICROS_PER_SEC,
+
std::memory_order_relaxed);
+}
+
void RpcRateLimiter::reset(int qps) {
limiter->reset(qps, qps, 0);
}
@@ -129,7 +137,8 @@ void HostLevelMSRpcRateLimiters::init_with_uniform_qps(int
qps) {
}
int64_t HostLevelMSRpcRateLimiters::limit(MetaServiceRPC rpc) {
- if (!config::enable_ms_rpc_host_level_rate_limit) {
+ const bool dry_run = config::enable_ms_rpc_host_level_rate_limit_dry_run;
+ if (!config::enable_ms_rpc_host_level_rate_limit && !dry_run) {
return 0;
}
@@ -139,10 +148,21 @@ int64_t HostLevelMSRpcRateLimiters::limit(MetaServiceRPC
rpc) {
}
auto limiter = _limiters[idx].load();
- if (limiter && limiter->limiter) {
- return limiter->limiter->add(1);
+ if (!limiter) {
+ return 0;
+ }
+ DCHECK(limiter->limiter);
+
+ auto result = dry_run ? limiter->limiter->reserve_with_config(1)
+ : limiter->limiter->add_with_config(1);
+ if (result.sleep_duration > 0 && limiter->should_log(MonotonicMicros())) {
+ LOG(INFO) << "[ms-throttle] host-level rate limiter "
+ << (dry_run ? "dry run would throttle" : "throttled") << "
MS RPC request"
+ << ", rpc=" << meta_service_rpc_display_name(rpc)
+ << (dry_run ? ", estimated_wait_ns=" : ", sleep_ns=") <<
result.sleep_duration
+ << ", qps_limit=" << result.max_speed;
}
- return 0;
+ return dry_run ? 0 : result.sleep_duration;
}
void HostLevelMSRpcRateLimiters::reset(MetaServiceRPC rpc, int qps) {
diff --git a/be/src/cloud/cloud_ms_rpc_rate_limiters.h
b/be/src/cloud/cloud_ms_rpc_rate_limiters.h
index 17a4bf465ab..e0190f208e7 100644
--- a/be/src/cloud/cloud_ms_rpc_rate_limiters.h
+++ b/be/src/cloud/cloud_ms_rpc_rate_limiters.h
@@ -20,6 +20,7 @@
#include <bvar/bvar.h>
#include <array>
+#include <atomic>
#include <memory>
#include <string>
#include <string_view>
@@ -78,8 +79,14 @@ struct RpcRateLimiter {
RpcRateLimiter(int qps, std::string_view op_name);
+ // Each RPC type owns one RpcRateLimiter, so log suppression is
independent per RPC type.
+ bool should_log(int64_t now_us);
+
// Reset the rate limiter with new QPS
void reset(int qps);
+
+private:
+ std::atomic<int64_t> _next_log_time_us {0};
};
// Host-level rate limiters for MS RPCs to prevent burst traffic
@@ -96,7 +103,8 @@ public:
~HostLevelMSRpcRateLimiters() = default;
- // Rate limit the specified RPC method, returns actual sleep time in
nanoseconds
+ // Rate limit the specified RPC method, returning the actual sleep time in
nanoseconds.
+ // Dry-run mode updates the limiter and its metrics without sleeping, and
returns 0.
// Thread-safe: each limiter handles its own synchronization
int64_t limit(MetaServiceRPC rpc);
diff --git a/be/src/cloud/cloud_throttle_state_machine.cpp
b/be/src/cloud/cloud_throttle_state_machine.cpp
index cd174ee7b5d..cca0112a7ed 100644
--- a/be/src/cloud/cloud_throttle_state_machine.cpp
+++ b/be/src/cloud/cloud_throttle_state_machine.cpp
@@ -83,23 +83,26 @@ std::vector<RpcThrottleAction>
RpcThrottleStateMachine::on_upgrade(
// Apply floor
new_limit = std::max(new_limit, floor_qps);
- // Only apply if it's actually limiting
- if (new_limit < snapshot.current_qps || old_limit > 0) {
- RpcThrottleAction action {
- .type = RpcThrottleAction::Type::SET_LIMIT,
- .rpc_type = snapshot.rpc_type,
- .table_id = snapshot.table_id,
- .qps_limit = new_limit,
- };
- actions.push_back(action);
- record.changes[key] = {old_limit, new_limit};
- _current_limits[key] = new_limit;
-
- LOG(INFO) << "[ms-throttle] upgrade: rpc=" <<
load_related_rpc_name(snapshot.rpc_type)
- << ", table_id=" << snapshot.table_id
- << ", current_qps=" << snapshot.current_qps << ",
old_limit=" << old_limit
- << ", new_limit=" << new_limit;
+ // An upgrade must make the effective limit stricter. For a table
without an
+ // existing limit, its current QPS is the baseline.
+ const double baseline = old_limit > 0 ? old_limit :
snapshot.current_qps;
+ if (new_limit >= baseline) {
+ continue;
}
+
+ RpcThrottleAction action {
+ .type = RpcThrottleAction::Type::SET_LIMIT,
+ .rpc_type = snapshot.rpc_type,
+ .table_id = snapshot.table_id,
+ .qps_limit = new_limit,
+ };
+ actions.push_back(action);
+ record.changes[key] = {old_limit, new_limit};
+ _current_limits[key] = new_limit;
+
+ LOG(INFO) << "[ms-throttle] upgrade: rpc=" <<
load_related_rpc_name(snapshot.rpc_type)
+ << ", table_id=" << snapshot.table_id << ", current_qps=" <<
snapshot.current_qps
+ << ", old_limit=" << old_limit << ", new_limit=" <<
new_limit;
}
if (!record.changes.empty()) {
@@ -137,6 +140,7 @@ std::vector<RpcThrottleAction>
RpcThrottleStateMachine::on_downgrade() {
.rpc_type = rpc_type,
.table_id = table_id,
.qps_limit = old_limit,
+ .reset_reservation = true,
};
actions.push_back(action);
@@ -189,6 +193,14 @@ RpcThrottleParams RpcThrottleStateMachine::get_params()
const {
// ============== RpcThrottleCoordinator ==============
+static void advance_tick_counter(int64_t& counter, int64_t ticks, int64_t
stop_at) {
+ DCHECK_GE(ticks, 0);
+ if (counter < 0 || counter >= stop_at) {
+ return;
+ }
+ counter += std::min(ticks, stop_at - counter);
+}
+
RpcThrottleCoordinator::RpcThrottleCoordinator(ThrottleCoordinatorParams
params) : _params(params) {
LOG(INFO) << "[ms-throttle] coordinator initialized:
upgrade_cooldown_ticks="
<< params.upgrade_cooldown_ticks
@@ -221,22 +233,29 @@ bool RpcThrottleCoordinator::report_ms_busy() {
<< ", cooldown=" << _params.upgrade_cooldown_ticks;
return true; // Should trigger upgrade
}
+
+ if (!_has_pending_upgrades) {
+ _ticks_since_last_ms_busy = -1;
+ }
return false; // Cooling down
}
-bool RpcThrottleCoordinator::tick(int ticks) {
+bool RpcThrottleCoordinator::tick(int64_t ticks) {
std::lock_guard lock(_mtx);
- // Increment tick counters
- if (_ticks_since_last_ms_busy >= 0) {
- _ticks_since_last_ms_busy += ticks;
- }
- if (_ticks_since_last_upgrade >= 0) {
- _ticks_since_last_upgrade += ticks;
+ // The upgrade counter is needed even without pending history to preserve
cooldown.
+ // Stop at the threshold because larger values do not change the decision.
+ advance_tick_counter(_ticks_since_last_upgrade, ticks,
_params.upgrade_cooldown_ticks);
+
+ if (!_has_pending_upgrades) {
+ _ticks_since_last_ms_busy = -1;
+ return false;
}
+ advance_tick_counter(_ticks_since_last_ms_busy, ticks,
_params.downgrade_after_ticks);
+
// Check if downgrade should be triggered
- if (_has_pending_upgrades && _ticks_since_last_ms_busy >=
_params.downgrade_after_ticks) {
+ if (_ticks_since_last_ms_busy >= _params.downgrade_after_ticks) {
// Reset for next downgrade cycle
auto actual_ticks = _ticks_since_last_ms_busy;
_ticks_since_last_ms_busy = 0;
@@ -252,14 +271,17 @@ bool RpcThrottleCoordinator::tick(int ticks) {
void RpcThrottleCoordinator::set_has_pending_upgrades(bool has) {
std::lock_guard lock(_mtx);
_has_pending_upgrades = has;
+ if (!has) {
+ _ticks_since_last_ms_busy = -1;
+ }
}
-int RpcThrottleCoordinator::ticks_since_last_ms_busy() const {
+int64_t RpcThrottleCoordinator::ticks_since_last_ms_busy() const {
std::lock_guard lock(_mtx);
return _ticks_since_last_ms_busy;
}
-int RpcThrottleCoordinator::ticks_since_last_upgrade() const {
+int64_t RpcThrottleCoordinator::ticks_since_last_upgrade() const {
std::lock_guard lock(_mtx);
return _ticks_since_last_upgrade;
}
diff --git a/be/src/cloud/cloud_throttle_state_machine.h
b/be/src/cloud/cloud_throttle_state_machine.h
index 9a5f53c079b..d85597e1513 100644
--- a/be/src/cloud/cloud_throttle_state_machine.h
+++ b/be/src/cloud/cloud_throttle_state_machine.h
@@ -58,6 +58,8 @@ struct RpcThrottleAction {
LoadRelatedRpc rpc_type;
int64_t table_id;
double qps_limit {0}; // only meaningful for SET_LIMIT
+ // Discard queued reservations when applying a downgraded SET_LIMIT action.
+ bool reset_reservation {false};
};
// ============== ThrottleStateMachine ==============
@@ -153,24 +155,26 @@ public:
// Advance by specified number of ticks (caller decides actual time
between ticks)
// Returns true if downgrade should be triggered
- bool tick(int ticks = 1);
+ bool tick(int64_t ticks = 1);
// Tell coordinator whether there are pending upgrades that can be
downgraded
// Called by the state machine consumer after upgrade/downgrade
void set_has_pending_upgrades(bool has);
// Query state
- int ticks_since_last_ms_busy() const;
- int ticks_since_last_upgrade() const;
+ int64_t ticks_since_last_ms_busy() const;
+ int64_t ticks_since_last_upgrade() const;
ThrottleCoordinatorParams get_params() const;
private:
mutable std::mutex _mtx;
ThrottleCoordinatorParams _params;
- int _ticks_since_last_ms_busy = -1; // -1 means never received
- int _ticks_since_last_upgrade = -1; // -1 means never upgraded
- bool _has_pending_upgrades = false; // Whether there are upgrade records
to downgrade
+ // Counters saturate at their decision thresholds. The MS_BUSY counter is
reset
+ // when there is no pending upgrade history, so neither counter grows
unbounded.
+ int64_t _ticks_since_last_ms_busy = -1; // -1 means inactive or never
received
+ int64_t _ticks_since_last_upgrade = -1; // -1 means never upgraded
+ bool _has_pending_upgrades = false; // Whether there are upgrade
records to downgrade
};
} // namespace doris::cloud
diff --git a/be/src/cloud/config.cpp b/be/src/cloud/config.cpp
index 8beeb1ba806..4f764a8ff92 100644
--- a/be/src/cloud/config.cpp
+++ b/be/src/cloud/config.cpp
@@ -192,6 +192,7 @@
DEFINE_mBool(enable_file_cache_write_cumu_compaction_index_only, "false");
// MS RPC rate limiting config
DEFINE_mBool(enable_ms_rpc_host_level_rate_limit, "false");
+DEFINE_mBool(enable_ms_rpc_host_level_rate_limit_dry_run, "true");
// Per-RPC QPS limit configs (per CPU core)
// QPS limit = config_value * num_cores
@@ -223,6 +224,7 @@ DEFINE_mInt32(ms_rpc_qps_update_packed_file_info, "-1");
// Table-level backpressure handling config
DEFINE_mBool(enable_ms_backpressure_handling, "false");
+DEFINE_mBool(enable_ms_backpressure_handling_dry_run, "true");
DEFINE_Int32(ms_rpc_table_qps_window_sec, "3");
// Throttle upgrade config
diff --git a/be/src/cloud/config.h b/be/src/cloud/config.h
index 15162b04dc0..aa271889344 100644
--- a/be/src/cloud/config.h
+++ b/be/src/cloud/config.h
@@ -233,6 +233,10 @@
DECLARE_mBool(enable_file_cache_write_cumu_compaction_index_only);
// MS RPC rate limiting config
// Enable host-level rate limiting for MS RPCs to prevent burst traffic
DECLARE_mBool(enable_ms_rpc_host_level_rate_limit);
+// Evaluate and record host-level MS RPC rate limits without delaying requests.
+// Dry-run evaluation is independent of enable_ms_rpc_host_level_rate_limit.
+// When both are enabled, dry-run takes precedence and requests are not
delayed.
+DECLARE_mBool(enable_ms_rpc_host_level_rate_limit_dry_run);
// Per-RPC QPS limit configs (per CPU core)
// QPS limit = config_value * num_cores
@@ -266,6 +270,10 @@ DECLARE_mInt32(ms_rpc_qps_update_packed_file_info);
// Enable MS backpressure response handling (table-level adaptive throttling)
DECLARE_mBool(enable_ms_backpressure_handling);
+// Evaluate and record table-level adaptive throttling without delaying
requests.
+// Dry-run evaluation is independent of enable_ms_backpressure_handling.
+// When both are enabled, dry-run takes precedence and requests are not
delayed.
+DECLARE_mBool(enable_ms_backpressure_handling_dry_run);
// Time window (seconds) for computing per-table QPS via bvar::PerSecond.
// Larger window smooths out short-term spikes; smaller window reacts faster.
diff --git a/be/test/cloud/cloud_ms_backpressure_handler_test.cpp
b/be/test/cloud/cloud_ms_backpressure_handler_test.cpp
index 4f391c54b5d..c3b4fcea368 100644
--- a/be/test/cloud/cloud_ms_backpressure_handler_test.cpp
+++ b/be/test/cloud/cloud_ms_backpressure_handler_test.cpp
@@ -26,6 +26,7 @@
#include <vector>
#include "cloud/config.h"
+#include "util/time.h"
namespace doris::cloud {
@@ -58,6 +59,28 @@ TEST_F(StrictQpsLimiterTest, UpdateQps) {
EXPECT_DOUBLE_EQ(limiter.get_qps(), 100.0);
}
+TEST_F(StrictQpsLimiterTest, UpdateQpsCanResetReservations) {
+ StrictQpsLimiter limiter(1.0);
+
+ limiter.reserve();
+ auto queued = limiter.reserve();
+ EXPECT_GT(queued, std::chrono::steady_clock::now());
+
+ limiter.update_qps(2.0, true);
+ auto now = std::chrono::steady_clock::now();
+ auto first_after_reset = limiter.reserve();
+ auto reset_delay_ms =
+
std::chrono::duration_cast<std::chrono::milliseconds>(first_after_reset -
now).count();
+ EXPECT_LE(reset_delay_ms, 10);
+
+ auto second_after_reset = limiter.reserve();
+ auto interval_ms =
std::chrono::duration_cast<std::chrono::milliseconds>(second_after_reset -
+
first_after_reset)
+ .count();
+ EXPECT_GE(interval_ms, 490);
+ EXPECT_LE(interval_ms, 510);
+}
+
TEST_F(StrictQpsLimiterTest, ZeroQpsDefaultsToOne) {
StrictQpsLimiter limiter(0.0);
EXPECT_DOUBLE_EQ(limiter.get_qps(), 1.0);
@@ -266,6 +289,56 @@ TEST_F(TableRpcQpsRegistryTest,
GetTopKTablesInvalidRpcType) {
EXPECT_TRUE(top_tables.empty());
}
+TEST_F(TableRpcQpsRegistryTest, CleanupKeepsActiveCounters) {
+ TableRpcQpsRegistry registry(std::chrono::hours(1),
std::chrono::minutes(1));
+ registry.record(LoadRelatedRpc::PREPARE_ROWSET, 100);
+
+ EXPECT_EQ(registry.cleanup_inactive_tables(), 0);
+
EXPECT_EQ(registry.get_tracked_table_count(LoadRelatedRpc::PREPARE_ROWSET), 1);
+}
+
+TEST_F(TableRpcQpsRegistryTest, CleanupRemovesInactiveCounters) {
+ TableRpcQpsRegistry registry(std::chrono::hours(1),
std::chrono::milliseconds(0));
+ registry.record(LoadRelatedRpc::PREPARE_ROWSET, 100);
+ registry.record(LoadRelatedRpc::COMMIT_ROWSET, 200);
+
+ EXPECT_EQ(registry.cleanup_inactive_tables(), 2);
+
EXPECT_EQ(registry.get_tracked_table_count(LoadRelatedRpc::PREPARE_ROWSET), 0);
+ EXPECT_EQ(registry.get_tracked_table_count(LoadRelatedRpc::COMMIT_ROWSET),
0);
+}
+
+TEST_F(TableRpcQpsRegistryTest, CleanupThreadRunsIndependently) {
+ TableRpcQpsRegistry registry(std::chrono::milliseconds(10),
std::chrono::milliseconds(50));
+ registry.record(LoadRelatedRpc::PREPARE_ROWSET, 100);
+
EXPECT_EQ(registry.get_tracked_table_count(LoadRelatedRpc::PREPARE_ROWSET), 1);
+
+ for (int i = 0;
+ i < 100 &&
registry.get_tracked_table_count(LoadRelatedRpc::PREPARE_ROWSET) != 0; ++i) {
+ std::this_thread::sleep_for(std::chrono::milliseconds(10));
+ }
+
EXPECT_EQ(registry.get_tracked_table_count(LoadRelatedRpc::PREPARE_ROWSET), 0);
+}
+
+TEST_F(TableRpcQpsRegistryTest, ConcurrentRecordAndCleanup) {
+ TableRpcQpsRegistry registry(std::chrono::hours(1),
std::chrono::milliseconds(0));
+
+ std::thread recorder([®istry]() {
+ for (int i = 0; i < 10000; ++i) {
+ registry.record(LoadRelatedRpc::PREPARE_ROWSET, i % 100);
+ }
+ });
+ std::thread cleaner([®istry]() {
+ for (int i = 0; i < 1000; ++i) {
+ registry.cleanup_inactive_tables();
+ }
+ });
+ recorder.join();
+ cleaner.join();
+
+ registry.record(LoadRelatedRpc::PREPARE_ROWSET, 100);
+
EXPECT_GE(registry.get_tracked_table_count(LoadRelatedRpc::PREPARE_ROWSET), 1);
+}
+
// ============== TableRpcThrottler Tests ==============
class TableRpcThrottlerTest : public testing::Test {
@@ -324,6 +397,56 @@ TEST_F(TableRpcThrottlerTest, ThrottleWithLimit) {
EXPECT_LE(diff_ms, 1100);
}
+TEST_F(TableRpcThrottlerTest, DryRunReservationIsSharedWithActualLimiter) {
+ TableRpcThrottler throttler;
+ throttler.set_qps_limit(LoadRelatedRpc::UPDATE_DELETE_BITMAP, 300, 1.0);
+
+ auto dry_run_t1 = throttler.throttle(LoadRelatedRpc::UPDATE_DELETE_BITMAP,
300, true);
+ auto dry_run_t2 = throttler.throttle(LoadRelatedRpc::UPDATE_DELETE_BITMAP,
300, true);
+ auto dry_run_diff_ms =
std::chrono::duration_cast<std::chrono::milliseconds>(
+ dry_run_t2.wait_until -
dry_run_t1.wait_until)
+ .count();
+ EXPECT_GE(dry_run_diff_ms, 900);
+ EXPECT_LE(dry_run_diff_ms, 1100);
+
+ auto actual_t1 = throttler.throttle(LoadRelatedRpc::UPDATE_DELETE_BITMAP,
300, false);
+ auto shared_diff_ms =
std::chrono::duration_cast<std::chrono::milliseconds>(
+ actual_t1.wait_until - dry_run_t2.wait_until)
+ .count();
+ EXPECT_GE(shared_diff_ms, 900);
+ EXPECT_LE(shared_diff_ms, 1100);
+}
+
+TEST_F(TableRpcThrottlerTest, DowngradeResetsReservations) {
+ TableRpcThrottler throttler;
+ throttler.set_qps_limit(LoadRelatedRpc::UPDATE_DELETE_BITMAP, 300, 1.0);
+
+ throttler.throttle(LoadRelatedRpc::UPDATE_DELETE_BITMAP, 300, true);
+ auto queued = throttler.throttle(LoadRelatedRpc::UPDATE_DELETE_BITMAP,
300, true);
+ EXPECT_GT(queued.wait_until, std::chrono::steady_clock::now());
+
+ throttler.set_qps_limit(LoadRelatedRpc::UPDATE_DELETE_BITMAP, 300, 2.0,
true);
+ auto now = std::chrono::steady_clock::now();
+ auto after_downgrade =
throttler.throttle(LoadRelatedRpc::UPDATE_DELETE_BITMAP, 300, false);
+ auto reset_delay_ms =
+
std::chrono::duration_cast<std::chrono::milliseconds>(after_downgrade.wait_until
- now)
+ .count();
+ EXPECT_LE(reset_delay_ms, 10);
+}
+
+TEST_F(TableRpcThrottlerTest, ThrottleLogIsRateLimitedPerRpc) {
+ TableRpcThrottler throttler;
+ constexpr int64_t first_log_time_us = 100;
+
+ EXPECT_TRUE(throttler.should_log(LoadRelatedRpc::PREPARE_ROWSET,
first_log_time_us));
+ EXPECT_FALSE(throttler.should_log(LoadRelatedRpc::PREPARE_ROWSET,
+ first_log_time_us + MICROS_PER_SEC - 1));
+ EXPECT_TRUE(throttler.should_log(LoadRelatedRpc::COMMIT_ROWSET,
+ first_log_time_us + MICROS_PER_SEC - 1));
+ EXPECT_TRUE(throttler.should_log(LoadRelatedRpc::PREPARE_ROWSET,
+ first_log_time_us + MICROS_PER_SEC));
+}
+
TEST_F(TableRpcThrottlerTest, ThrottledTableCount) {
TableRpcThrottler throttler;
@@ -355,6 +478,8 @@ class MSBackpressureHandlerTest : public testing::Test {
protected:
void SetUp() override {
_saved_enable = config::enable_ms_backpressure_handling;
+ _saved_dry_run = config::enable_ms_backpressure_handling_dry_run;
+ config::enable_ms_backpressure_handling_dry_run = false;
_saved_upgrade_interval = config::ms_backpressure_upgrade_interval_ms;
_saved_downgrade_interval =
config::ms_backpressure_downgrade_interval_ms;
_saved_top_k = config::ms_backpressure_upgrade_top_k;
@@ -364,6 +489,7 @@ protected:
void TearDown() override {
config::enable_ms_backpressure_handling = _saved_enable;
+ config::enable_ms_backpressure_handling_dry_run = _saved_dry_run;
config::ms_backpressure_upgrade_interval_ms = _saved_upgrade_interval;
config::ms_backpressure_downgrade_interval_ms =
_saved_downgrade_interval;
config::ms_backpressure_upgrade_top_k = _saved_top_k;
@@ -373,6 +499,7 @@ protected:
private:
bool _saved_enable;
+ bool _saved_dry_run;
int32_t _saved_upgrade_interval;
int32_t _saved_downgrade_interval;
int32_t _saved_top_k;
@@ -380,8 +507,9 @@ private:
double _saved_floor;
};
-TEST_F(MSBackpressureHandlerTest, DisabledByDefault) {
+TEST_F(MSBackpressureHandlerTest, DisabledWhenActualAndDryRunAreOff) {
config::enable_ms_backpressure_handling = false;
+ config::enable_ms_backpressure_handling_dry_run = false;
TableRpcQpsRegistry registry;
TableRpcThrottler throttler;
@@ -438,10 +566,12 @@ TEST_F(MSBackpressureHandlerTest, BeforeAndAfterRpc) {
// before_rpc with no limit should return approximately now
auto now = std::chrono::steady_clock::now();
- auto wait_until = handler.before_rpc(LoadRelatedRpc::COMMIT_ROWSET, 12345);
+ auto decision = handler.before_rpc(LoadRelatedRpc::COMMIT_ROWSET, 12345);
- auto diff =
std::chrono::duration_cast<std::chrono::milliseconds>(wait_until - now).count();
+ auto diff =
std::chrono::duration_cast<std::chrono::milliseconds>(decision.wait_until - now)
+ .count();
EXPECT_LE(diff, 10);
+ EXPECT_FALSE(decision.dry_run);
// after_rpc should record the call (just verify it doesn't crash)
handler.after_rpc(LoadRelatedRpc::COMMIT_ROWSET, 12345);
@@ -463,9 +593,39 @@ TEST_F(MSBackpressureHandlerTest, BeforeRpcWithThrottle) {
// Second call should return a time ~1 second later
auto t2 = handler.before_rpc(LoadRelatedRpc::UPDATE_DELETE_BITMAP, 500);
- auto diff_ms = std::chrono::duration_cast<std::chrono::milliseconds>(t2 -
t1).count();
+ auto diff_ms =
+
std::chrono::duration_cast<std::chrono::milliseconds>(t2.wait_until -
t1.wait_until)
+ .count();
+ EXPECT_GE(diff_ms, 900);
+ EXPECT_LE(diff_ms, 1100);
+ EXPECT_FALSE(t2.dry_run);
+ EXPECT_DOUBLE_EQ(t2.qps_limit, 1.0);
+}
+
+TEST_F(MSBackpressureHandlerTest, DryRunWorksWithoutActualEnforcement) {
+ config::enable_ms_backpressure_handling = false;
+ config::enable_ms_backpressure_handling_dry_run = true;
+ config::ms_backpressure_upgrade_interval_ms = 0;
+
+ TableRpcQpsRegistry registry;
+ TableRpcThrottler throttler;
+ MSBackpressureHandler handler(®istry, &throttler);
+
+ EXPECT_TRUE(handler.on_ms_busy());
+
+ throttler.set_qps_limit(LoadRelatedRpc::UPDATE_DELETE_BITMAP, 500, 1.0);
+ auto t1 = handler.before_rpc(LoadRelatedRpc::UPDATE_DELETE_BITMAP, 500);
+ auto t2 = handler.before_rpc(LoadRelatedRpc::UPDATE_DELETE_BITMAP, 500);
+
+ auto diff_ms =
+
std::chrono::duration_cast<std::chrono::milliseconds>(t2.wait_until -
t1.wait_until)
+ .count();
EXPECT_GE(diff_ms, 900);
EXPECT_LE(diff_ms, 1100);
+ EXPECT_TRUE(t2.dry_run);
+ EXPECT_DOUBLE_EQ(t2.qps_limit, 1.0);
+
+ handler.after_rpc(LoadRelatedRpc::UPDATE_DELETE_BITMAP, 500);
}
TEST_F(MSBackpressureHandlerTest, SecondsSinceLastMsBusy) {
diff --git a/be/test/cloud/cloud_ms_rpc_rate_limiters_test.cpp
b/be/test/cloud/cloud_ms_rpc_rate_limiters_test.cpp
index 27a141e87ea..82b60cd15f0 100644
--- a/be/test/cloud/cloud_ms_rpc_rate_limiters_test.cpp
+++ b/be/test/cloud/cloud_ms_rpc_rate_limiters_test.cpp
@@ -20,23 +20,33 @@
#include <gtest/gtest.h>
#include <atomic>
+#include <chrono>
#include <thread>
#include <vector>
#include "cloud/config.h"
#include "util/cpu_info.h"
+#include "util/time.h"
namespace doris::cloud {
// Basic tests using uniform QPS constructor (completely independent of config
and CPU cores)
class HostLevelMSRpcRateLimitersTest : public testing::Test {
protected:
- void SetUp() override { _saved_enable =
config::enable_ms_rpc_host_level_rate_limit; }
+ void SetUp() override {
+ _saved_enable = config::enable_ms_rpc_host_level_rate_limit;
+ _saved_dry_run = config::enable_ms_rpc_host_level_rate_limit_dry_run;
+ config::enable_ms_rpc_host_level_rate_limit_dry_run = false;
+ }
- void TearDown() override { config::enable_ms_rpc_host_level_rate_limit =
_saved_enable; }
+ void TearDown() override {
+ config::enable_ms_rpc_host_level_rate_limit = _saved_enable;
+ config::enable_ms_rpc_host_level_rate_limit_dry_run = _saved_dry_run;
+ }
private:
bool _saved_enable;
+ bool _saved_dry_run;
};
// Test that limit returns 0 when rate limiting is disabled
@@ -81,6 +91,55 @@ TEST_F(HostLevelMSRpcRateLimitersTest,
RateLimitingThrottles) {
EXPECT_GT(total_sleep_ns, 0) << "Rate limiting should have caused some
sleep";
}
+// Test that dry-run mode evaluates throttling without delaying requests,
regardless of the
+// enforcement switch.
+TEST_F(HostLevelMSRpcRateLimitersTest, DryRunObservesWithoutThrottling) {
+ config::enable_ms_rpc_host_level_rate_limit_dry_run = true;
+
+ for (bool rate_limit_enabled : {false, true}) {
+ config::enable_ms_rpc_host_level_rate_limit = rate_limit_enabled;
+ HostLevelMSRpcRateLimiters limiters(1);
+
+ EXPECT_EQ(limiters.limit(MetaServiceRPC::GET_TABLET_META), 0);
+ auto start = std::chrono::steady_clock::now();
+ for (int i = 0; i < 4; ++i) {
+ EXPECT_EQ(limiters.limit(MetaServiceRPC::GET_TABLET_META), 0);
+ }
+ auto elapsed = std::chrono::steady_clock::now() - start;
+
+ size_t idx = static_cast<size_t>(MetaServiceRPC::GET_TABLET_META);
+ auto limiter = limiters._limiters[idx].load();
+ ASSERT_NE(limiter, nullptr);
+ EXPECT_EQ(limiter->latency_recorder->count(), 4);
+ EXPECT_LT(elapsed, std::chrono::seconds(2));
+ }
+}
+
+TEST_F(HostLevelMSRpcRateLimitersTest,
DryRunReservationIsSharedWithEnforcement) {
+ config::enable_ms_rpc_host_level_rate_limit = false;
+ config::enable_ms_rpc_host_level_rate_limit_dry_run = true;
+ HostLevelMSRpcRateLimiters limiters(10);
+
+ for (int i = 0; i < 11; ++i) {
+ EXPECT_EQ(limiters.limit(MetaServiceRPC::GET_TABLET_META), 0);
+ }
+
+ config::enable_ms_rpc_host_level_rate_limit = true;
+ config::enable_ms_rpc_host_level_rate_limit_dry_run = false;
+ EXPECT_GT(limiters.limit(MetaServiceRPC::GET_TABLET_META), 0);
+}
+
+TEST_F(HostLevelMSRpcRateLimitersTest, RateLimitLogIsRateLimitedPerRpc) {
+ RpcRateLimiter get_tablet_meta_limiter(1, "rate limit log get tablet meta
test");
+ RpcRateLimiter get_rowset_limiter(1, "rate limit log get rowset test");
+ constexpr int64_t first_log_time_us = 100;
+
+ EXPECT_TRUE(get_tablet_meta_limiter.should_log(first_log_time_us));
+ EXPECT_FALSE(get_tablet_meta_limiter.should_log(first_log_time_us +
MICROS_PER_SEC - 1));
+ EXPECT_TRUE(get_rowset_limiter.should_log(first_log_time_us +
MICROS_PER_SEC - 1));
+ EXPECT_TRUE(get_tablet_meta_limiter.should_log(first_log_time_us +
MICROS_PER_SEC));
+}
+
// Test multiple RPC types have independent rate limiters
TEST_F(HostLevelMSRpcRateLimitersTest, IndependentRateLimiters) {
config::enable_ms_rpc_host_level_rate_limit = true;
diff --git a/be/test/cloud/cloud_throttle_state_machine_test.cpp
b/be/test/cloud/cloud_throttle_state_machine_test.cpp
index d5fc70c8647..cb7ef37edeb 100644
--- a/be/test/cloud/cloud_throttle_state_machine_test.cpp
+++ b/be/test/cloud/cloud_throttle_state_machine_test.cpp
@@ -19,6 +19,8 @@
#include <gtest/gtest.h>
+#include <limits>
+
namespace doris::cloud {
// ============== RpcThrottleStateMachine Tests ==============
@@ -76,6 +78,7 @@ TEST_F(RpcThrottleStateMachineTest,
MultipleUpgradesThenDowngrades) {
auto a1 = sm.on_upgrade({{LoadRelatedRpc::COMMIT_ROWSET, 100, 80.0}});
ASSERT_EQ(a1.size(), 1);
EXPECT_DOUBLE_EQ(a1[0].qps_limit, 40.0); // 80 * 0.5
+ EXPECT_FALSE(a1[0].reset_reservation);
// Second upgrade, same table, current limit is 40
auto a2 = sm.on_upgrade({{LoadRelatedRpc::COMMIT_ROWSET, 100, 40.0}});
@@ -89,6 +92,7 @@ TEST_F(RpcThrottleStateMachineTest,
MultipleUpgradesThenDowngrades) {
ASSERT_EQ(d1.size(), 1);
EXPECT_EQ(d1[0].type, RpcThrottleAction::Type::SET_LIMIT);
EXPECT_DOUBLE_EQ(d1[0].qps_limit, 40.0);
+ EXPECT_TRUE(d1[0].reset_reservation);
// Second downgrade: undo first upgrade, remove limit
auto d2 = sm.on_downgrade();
@@ -170,9 +174,15 @@ TEST_F(RpcThrottleStateMachineTest,
UpdateFloorQpsAtRuntime) {
// Runtime update floor_qps=5.0
sm.update_params({.top_k = 1, .ratio = 0.01, .floor_qps = 5.0});
- // Second upgrade, new floor takes effect
- auto a2 = sm.on_upgrade({{LoadRelatedRpc::PREPARE_ROWSET, 100, 1.0}});
- EXPECT_DOUBLE_EQ(a2[0].qps_limit, 5.0); // 1*0.01=0.01 < floor(5.0)
+ // The higher floor applies to new limits but must not relax the existing
limit.
+ auto a2 = sm.on_upgrade({
+ {LoadRelatedRpc::PREPARE_ROWSET, 100, 1.0},
+ {LoadRelatedRpc::PREPARE_ROWSET, 200, 10.0},
+ });
+ ASSERT_EQ(a2.size(), 1);
+ EXPECT_EQ(a2[0].table_id, 200);
+ EXPECT_DOUBLE_EQ(a2[0].qps_limit, 5.0);
+ EXPECT_DOUBLE_EQ(sm.get_current_limit(LoadRelatedRpc::PREPARE_ROWSET,
100), 1.0);
}
TEST_F(RpcThrottleStateMachineTest, MultipleRpcTypes) {
@@ -436,9 +446,20 @@ TEST_F(RpcThrottleStateMachineTest,
FloorQpsWithRepeatedUpgrades) {
auto a4 = sm.on_upgrade({{LoadRelatedRpc::PREPARE_ROWSET, 100, 12.5}});
EXPECT_DOUBLE_EQ(a4[0].qps_limit, 10.0);
- // Already at floor: 10 * 0.5 = 5 < floor(10), stays at floor
- auto a5 = sm.on_upgrade({{LoadRelatedRpc::PREPARE_ROWSET, 100, 10.0}});
- EXPECT_DOUBLE_EQ(a5[0].qps_limit, 10.0);
+ EXPECT_EQ(sm.upgrade_level(), 4);
+
+ // Repeated upgrades at the floor are no-ops and do not add rollback
history.
+ for (int i = 0; i < 3; ++i) {
+ auto no_op = sm.on_upgrade({{LoadRelatedRpc::PREPARE_ROWSET, 100,
10.0}});
+ EXPECT_TRUE(no_op.empty());
+ EXPECT_EQ(sm.upgrade_level(), 4);
+ }
+
+ // One downgrade immediately restores the limit before reaching the floor.
+ auto downgrade = sm.on_downgrade();
+ ASSERT_EQ(downgrade.size(), 1);
+ EXPECT_DOUBLE_EQ(downgrade[0].qps_limit, 12.5);
+ EXPECT_EQ(sm.upgrade_level(), 3);
}
TEST_F(RpcThrottleStateMachineTest, MultiRpcTypeTopKIndependence) {
@@ -539,10 +560,66 @@ TEST_F(RpcThrottleCoordinatorTest,
NoDowngradeWithoutPendingUpgrades) {
// Explicitly clear pending upgrades to simulate the case where
// the caller decided not to upgrade (report_ms_busy sets it internally)
coord.set_has_pending_upgrades(false);
+ EXPECT_EQ(coord.ticks_since_last_ms_busy(), -1);
for (int i = 0; i < 100; i++) {
EXPECT_FALSE(coord.tick());
}
+ EXPECT_EQ(coord.ticks_since_last_ms_busy(), -1);
+ EXPECT_EQ(coord.ticks_since_last_upgrade(), params.upgrade_cooldown_ticks);
+}
+
+TEST_F(RpcThrottleCoordinatorTest,
MsBusyDuringIdleCooldownKeepsBusyCounterInactive) {
+ ThrottleCoordinatorParams params {.upgrade_cooldown_ticks = 10,
.downgrade_after_ticks = 3};
+ RpcThrottleCoordinator coord(params);
+
+ EXPECT_TRUE(coord.report_ms_busy());
+ coord.set_has_pending_upgrades(false);
+ EXPECT_FALSE(coord.tick());
+
+ EXPECT_FALSE(coord.report_ms_busy());
+ EXPECT_EQ(coord.ticks_since_last_ms_busy(), -1);
+}
+
+TEST_F(RpcThrottleCoordinatorTest, LargeTickSaturatesWithoutOverflow) {
+ ThrottleCoordinatorParams params {.upgrade_cooldown_ticks = 10,
.downgrade_after_ticks = 20};
+ RpcThrottleCoordinator coord(params);
+ constexpr int64_t kBeyondInt32 =
+ static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1000;
+
+ EXPECT_TRUE(coord.report_ms_busy());
+ coord.set_has_pending_upgrades(true);
+
+ EXPECT_TRUE(coord.tick(kBeyondInt32));
+ EXPECT_EQ(coord.ticks_since_last_ms_busy(), 0);
+ EXPECT_EQ(coord.ticks_since_last_upgrade(), params.upgrade_cooldown_ticks);
+
+ coord.set_has_pending_upgrades(false);
+ EXPECT_EQ(coord.ticks_since_last_ms_busy(), -1);
+ EXPECT_FALSE(coord.tick(kBeyondInt32));
+ EXPECT_EQ(coord.ticks_since_last_ms_busy(), -1);
+ EXPECT_EQ(coord.ticks_since_last_upgrade(), params.upgrade_cooldown_ticks);
+
+ // Saturating the upgrade counter preserves the cooldown decision while
idle.
+ EXPECT_TRUE(coord.report_ms_busy());
+}
+
+TEST_F(RpcThrottleCoordinatorTest,
IncreasedCooldownContinuesFromSaturatedCounter) {
+ ThrottleCoordinatorParams params {.upgrade_cooldown_ticks = 5,
.downgrade_after_ticks = 20};
+ RpcThrottleCoordinator coord(params);
+
+ EXPECT_TRUE(coord.report_ms_busy());
+ coord.set_has_pending_upgrades(false);
+ EXPECT_FALSE(coord.tick(5));
+ EXPECT_EQ(coord.ticks_since_last_upgrade(), 5);
+
+ coord.update_params({.upgrade_cooldown_ticks = 10, .downgrade_after_ticks
= 20});
+ EXPECT_FALSE(coord.tick(4));
+ EXPECT_EQ(coord.ticks_since_last_upgrade(), 9);
+ EXPECT_FALSE(coord.report_ms_busy());
+
+ EXPECT_FALSE(coord.tick());
+ EXPECT_TRUE(coord.report_ms_busy());
}
TEST_F(RpcThrottleCoordinatorTest, UpdateUpgradeCooldownAtRuntime) {
diff --git a/common/cpp/token_bucket_rate_limiter.cpp
b/common/cpp/token_bucket_rate_limiter.cpp
index 666fba7dbf1..5acf55d8435 100644
--- a/common/cpp/token_bucket_rate_limiter.cpp
+++ b/common/cpp/token_bucket_rate_limiter.cpp
@@ -112,6 +112,14 @@ std::pair<size_t, double>
TokenBucketRateLimiter::_update_remain_token(long now,
}
int64_t TokenBucketRateLimiter::add(size_t amount) {
+ int64_t sleep_time_ns = reserve(amount);
+ if (sleep_time_ns > 0) {
+ bthread_usleep(sleep_time_ns / 1000);
+ }
+ return sleep_time_ns;
+}
+
+int64_t TokenBucketRateLimiter::reserve(size_t amount) {
// Values obtained under lock to be checked after release
auto duration = std::chrono::steady_clock::now().time_since_epoch();
auto time_nano_count =
std::chrono::duration_cast<std::chrono::nanoseconds>(duration).count();
@@ -126,7 +134,6 @@ int64_t TokenBucketRateLimiter::add(size_t amount) {
int64_t sleep_time_ns = 0;
if (_max_speed && tokens_value < 0) {
sleep_time_ns = static_cast<int64_t>(-tokens_value / _max_speed * NS);
- bthread_usleep(sleep_time_ns / 1000);
}
return sleep_time_ns;
@@ -152,15 +159,25 @@ int64_t TokenBucketRateLimiterHolder::add(size_t amount) {
}
TokenBucketRateLimiterResult
TokenBucketRateLimiterHolder::add_with_config(size_t amount) {
- // Snapshot the current limiter and call add() outside the read lock:
add() may
- // sleep for a long time when throttled, and holding the read lock across
the
- // sleep would block reset() (dynamic config update) for the whole
duration.
+ return _consume_with_config(amount, true);
+}
+
+TokenBucketRateLimiterResult
TokenBucketRateLimiterHolder::reserve_with_config(size_t amount) {
+ return _consume_with_config(amount, false);
+}
+
+TokenBucketRateLimiterResult
TokenBucketRateLimiterHolder::_consume_with_config(size_t amount,
+
bool wait) {
+ // Snapshot the current limiter and consume outside the read lock. The
waiting
+ // path may sleep for a long time when throttled, and holding the read
lock across
+ // the sleep would block reset() (dynamic config update) for the whole
duration.
std::shared_ptr<TokenBucketRateLimiter> limiter;
{
std::shared_lock read {rate_limiter_rw_lock};
limiter = rate_limiter;
}
- TokenBucketRateLimiterResult result = {.sleep_duration =
limiter->add(amount),
+ TokenBucketRateLimiterResult result = {.sleep_duration = wait ?
limiter->add(amount)
+ :
limiter->reserve(amount),
.max_speed =
limiter->get_max_speed(),
.max_burst =
limiter->get_max_burst(),
.limit = limiter->get_limit()};
diff --git a/common/cpp/token_bucket_rate_limiter.h
b/common/cpp/token_bucket_rate_limiter.h
index 77eed4929e4..3c55336ef3a 100644
--- a/common/cpp/token_bucket_rate_limiter.h
+++ b/common/cpp/token_bucket_rate_limiter.h
@@ -59,6 +59,10 @@ public:
// Returns the sleep duration in nanoseconds, or -1 when the count limit
rejects the add.
int64_t add(size_t amount);
+ // Reserve `amount` tokens and return the required sleep duration without
sleeping.
+ // The token bucket state is updated in the same way as add().
+ int64_t reserve(size_t amount);
+
// Return `amount` tokens to the bucket (capped at max_burst) and roll
back the
// cumulative counter. Used to reconcile a reservation with the actually
consumed
// amount, e.g. a short read at EOF.
@@ -98,6 +102,8 @@ public:
int64_t add(size_t amount);
TokenBucketRateLimiterResult add_with_config(size_t amount);
+ // Reserve on the same limiter state as add_with_config(), but do not
sleep.
+ TokenBucketRateLimiterResult reserve_with_config(size_t amount);
// Charge `amount` like add(), but return the limiter generation the
tokens were
// taken from, or nullptr when the count limit rejects the charge. Callers
that later
@@ -117,6 +123,8 @@ public:
size_t get_limit() const;
private:
+ TokenBucketRateLimiterResult _consume_with_config(size_t amount, bool
wait);
+
mutable std::shared_mutex rate_limiter_rw_lock;
std::shared_ptr<TokenBucketRateLimiter> rate_limiter;
std::atomic<bool> _enabled;
diff --git
a/regression-test/suites/fault_injection_p0/cloud/test_cloud_ms_rpc_table_throttle.groovy
b/regression-test/suites/fault_injection_p0/cloud/test_cloud_ms_rpc_table_throttle.groovy
index 2293cf9709d..6119ce3e0e5 100644
---
a/regression-test/suites/fault_injection_p0/cloud/test_cloud_ms_rpc_table_throttle.groovy
+++
b/regression-test/suites/fault_injection_p0/cloud/test_cloud_ms_rpc_table_throttle.groovy
@@ -28,6 +28,7 @@ suite('test_cloud_ms_rpc_table_throttle', 'docker') {
options.setBeNum(1)
options.beConfigs += [
'enable_ms_backpressure_handling=true',
+ 'enable_ms_backpressure_handling_dry_run=false',
// Short intervals for faster test feedback
'ms_backpressure_upgrade_interval_ms=2000',
'ms_backpressure_downgrade_interval_ms=5000',
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
