Skip to content

[improvement](be) Add CPU-aware S3 rate limiter #65420

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.

Sign up for GitHub

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

Open
0AyanamiRei wants to merge 14 commits into
base: master
Choose a base branch
from
Open

[improvement](be) Add CPU-aware S3 rate limiter #65420

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
c2ffc30
[feature](be) Add CPU-aware S3 rate limiter configurations
0AyanamiRei Jul 9, 2026
a4bc63e
Rename token_per_second_max to s3_token_per_second_max for consistency
0AyanamiRei Jul 9, 2026
90a3236
[feature](be) Centralize S3 rate limiter config calculation
0AyanamiRei Jul 9, 2026
e48dfcf
[refactor](be) Use effective config for S3 rate limiter initializatio…
0AyanamiRei Jul 9, 2026
d529f3f
[test](be) Add unit tests for get_effective_s3_rate_limiter_config
0AyanamiRei Jul 9, 2026
e1eb6f9
refactor: Use S3RateLimiterConfigGuard in s3_rate_limiter test
0AyanamiRei Jul 9, 2026
0a221da
[test](be) Add unit tests for S3 rate limiter dynamic rebuild with pe…
0AyanamiRei Jul 9, 2026
1a1d991
[test](be) Assert get_limit() in S3 GET limiter cap-change rebuild test
0AyanamiRei Jul 9, 2026
85ba22a
[test](be) Fix stale config comment and add absolute PUT mode test
0AyanamiRei Jul 9, 2026
a4effbe
[fix](be) Guard CPU-aware S3 limiter overflow
0AyanamiRei Jul 9, 2026
46d20ea
remove
0AyanamiRei Jul 9, 2026
c614307
add
0AyanamiRei Jul 10, 2026
f6a3d0f
[improvement](be) Align S3 bandwidth limiter handling
0AyanamiRei Jul 10, 2026
72fcefa
rw
0AyanamiRei Jul 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions be/src/common/config.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1529,6 +1529,32 @@ DEFINE_Validator(s3_put_token_per_second, [](int64_t config) -> bool { return co

DEFINE_mInt64(s3_put_token_limit, "0");

// CPU-aware S3 rate limiter: QPS per CPU core. 0 means use absolute configs instead.
DEFINE_mInt64(s3_get_token_per_second_per_core, "0");
DEFINE_Validator(s3_get_token_per_second_per_core,
[](int64_t config) -> bool { return config >= 0; });

DEFINE_mInt64(s3_put_token_per_second_per_core, "0");
DEFINE_Validator(s3_put_token_per_second_per_core,
[](int64_t config) -> bool { return config >= 0; });

// Hard QPS cap for CPU-derived S3 rate limiters. 0 means no cap.
DEFINE_mInt64(s3_token_per_second_max, "0");
DEFINE_Validator(s3_token_per_second_max, [](int64_t config) -> bool { return config >= 0; });

// CPU-aware S3 bandwidth limiter: bytes per second per CPU core.
DEFINE_mInt64(s3_get_bytes_per_second_per_core, "104857600");
DEFINE_Validator(s3_get_bytes_per_second_per_core,
[](int64_t config) -> bool { return config >= 0; });

DEFINE_mInt64(s3_put_bytes_per_second_per_core, "104857600");
DEFINE_Validator(s3_put_bytes_per_second_per_core,
[](int64_t config) -> bool { return config >= 0; });

// Hard bandwidth cap for CPU-derived S3 bandwidth limiters. 0 means no cap.
DEFINE_mInt64(s3_bytes_per_second_max, "0");
DEFINE_Validator(s3_bytes_per_second_max, [](int64_t config) -> bool { return config >= 0; });

DEFINE_String(trino_connector_plugin_dir, "${DORIS_HOME}/plugins/connectors");

// ca_cert_file is in this path by default, Normally no modification is required
Expand Down
11 changes: 11 additions & 0 deletions be/src/common/config.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1609,6 +1609,17 @@ DECLARE_mInt64(s3_get_token_limit);
DECLARE_mInt64(s3_put_bucket_tokens);
DECLARE_mInt64(s3_put_token_per_second);
DECLARE_mInt64(s3_put_token_limit);

DECLARE_mInt64(s3_get_token_per_second_per_core);
DECLARE_mInt64(s3_put_token_per_second_per_core);
// optional
DECLARE_mInt64(s3_token_per_second_max);

DECLARE_mInt64(s3_get_bytes_per_second_per_core);
DECLARE_mInt64(s3_put_bytes_per_second_per_core);
// optional
DECLARE_mInt64(s3_bytes_per_second_max);

// max s3 client retry times
DECLARE_mInt32(max_s3_client_retry);
// When meet s3 429 error, the "get" request will
Expand Down
24 changes: 15 additions & 9 deletions be/src/io/fs/s3_obj_storage_client.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ inline ::Aws::Client::AWSError<::Aws::S3::S3Errors> s3_error_factory() {
}

template <typename Func>
auto s3_rate_limit(doris::S3RateLimitType op, Func callback) -> decltype(callback()) {
auto s3_rate_limit(doris::S3RateLimitType op, size_t bytes, Func callback) -> decltype(callback()) {
using T = decltype(callback());
if (!doris::config::enable_s3_rate_limiter) {
return callback();
Expand All @@ -86,17 +86,23 @@ auto s3_rate_limit(doris::S3RateLimitType op, Func callback) -> decltype(callbac
if (sleep_duration < 0) {
return T(s3_error_factory());
}
if (bytes > 0) {

@github-actions github-actions Bot Jul 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

The byte limiter only gets charged on the native S3 path. Azure still uses the same enable_s3_rate_limiter/S3ClientFactory::rate_limiter() helper, but its put_object, upload_part, and get_object calls never pass stream.size() or bytes_read through bytes_rate_limiter(). With the new positive defaults for s3_get_bytes_per_second_per_core and s3_put_bytes_per_second_per_core, enabling the rate limiter now throttles S3 bandwidth but silently leaves Azure bandwidth unlimited. Please either apply the same byte accounting to the Azure read/write helpers or make these new configs/API explicitly S3-native so Azure users do not expect them to apply.

sleep_duration = doris::S3ClientFactory::instance().bytes_rate_limiter(op)->add(bytes);
if (sleep_duration < 0) {
return T(s3_error_factory());
}
}
return callback();
}

template <typename Func>
auto s3_get_rate_limit(Func callback) -> decltype(callback()) {
return s3_rate_limit(doris::S3RateLimitType::GET, std::move(callback));
auto s3_get_rate_limit(Func callback, size_t bytes = 0) -> decltype(callback()) {
return s3_rate_limit(doris::S3RateLimitType::GET, bytes, std::move(callback));
}

template <typename Func>
auto s3_put_rate_limit(Func callback) -> decltype(callback()) {
return s3_rate_limit(doris::S3RateLimitType::PUT, std::move(callback));
auto s3_put_rate_limit(Func callback, size_t bytes = 0) -> decltype(callback()) {
return s3_rate_limit(doris::S3RateLimitType::PUT, bytes, std::move(callback));
}
} // namespace

Expand Down Expand Up @@ -167,7 +173,7 @@ ObjectStorageResponse S3ObjStorageClient::put_object(const ObjectStoragePathOpti
MonotonicStopWatch watch;
watch.start();
auto outcome = SYNC_POINT_HOOK_RETURN_VALUE(
s3_put_rate_limit([&]() { return _client->PutObject(request); }),
s3_put_rate_limit([&]() { return _client->PutObject(request); }, stream.size()),
"s3_file_writer::put_object", std::cref(request).get(), &stream);

watch.stop();
Expand Down Expand Up @@ -211,7 +217,7 @@ ObjectStorageUploadResponse S3ObjStorageClient::upload_part(const ObjectStorageP
MonotonicStopWatch watch;
watch.start();
auto outcome = SYNC_POINT_HOOK_RETURN_VALUE(
s3_put_rate_limit([&]() { return _client->UploadPart(request); }),
s3_put_rate_limit([&]() { return _client->UploadPart(request); }, stream.size()),
"s3_file_writer::upload_part", std::cref(request).get(), &stream);

watch.stop();
Expand Down Expand Up @@ -322,7 +328,7 @@ ObjectStorageResponse S3ObjStorageClient::get_object(const ObjectStoragePathOpti
request.SetResponseStreamFactory(AwsWriteableStreamFactory(buffer, bytes_read));

SCOPED_BVAR_LATENCY(s3_bvar::s3_get_latency);
auto outcome = s3_get_rate_limit([&]() { return _client->GetObject(request); });
auto outcome = s3_get_rate_limit([&]() { return _client->GetObject(request); }, bytes_read);
if (!outcome.IsSuccess()) {
return {convert_to_obj_response(s3fs_error(
outcome.GetError(), fmt::format("failed to read from {}", opts.key))),
Expand Down Expand Up @@ -502,4 +508,4 @@ std::string S3ObjStorageClient::generate_presigned_url(const ObjectStoragePathOp
expiration_secs);
}

} // namespace doris::io
} // namespace doris::io
180 changes: 170 additions & 10 deletions be/src/util/s3_util.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,9 @@
#include <cpp/token_bucket_rate_limiter.h>

#include <atomic>
#include <limits>

#include "util/cpu_info.h"
#include "util/string_util.h"

#ifdef USE_AZURE
Expand Down Expand Up @@ -156,28 +158,119 @@ constexpr char S3_NEED_OVERRIDE_ENDPOINT[] = "AWS_NEED_OVERRIDE_ENDPOINT";
constexpr char S3_ROLE_ARN[] = "AWS_ROLE_ARN";
constexpr char S3_EXTERNAL_ID[] = "AWS_EXTERNAL_ID";
constexpr char S3_CREDENTIALS_PROVIDER_TYPE[] = "AWS_CREDENTIALS_PROVIDER_TYPE";

static int64_t compute_s3_token_per_second_from_core(int64_t per_core, int64_t cores) {
const int64_t max_qps = config::s3_token_per_second_max > 0
? config::s3_token_per_second_max
: std::numeric_limits<int64_t>::max();
if (per_core > max_qps / cores) {
return max_qps;
}
return per_core * cores;
}

static int64_t compute_s3_bytes_per_second_from_core(int64_t per_core, int64_t cores) {
const int64_t max_bytes_per_second = config::s3_bytes_per_second_max > 0
? config::s3_bytes_per_second_max
: std::numeric_limits<int64_t>::max();
if (per_core > max_bytes_per_second / cores) {
return max_bytes_per_second;
}
return per_core * cores;
}
} // namespace

bvar::Adder<int64_t> get_rate_limit_ns("get_rate_limit_ns");
bvar::Adder<int64_t> get_rate_limit_exceed_req_num("get_rate_limit_exceed_req_num");
bvar::Adder<int64_t> put_rate_limit_ns("put_rate_limit_ns");
bvar::Adder<int64_t> put_rate_limit_exceed_req_num("put_rate_limit_exceed_req_num");
bvar::Adder<int64_t> get_bytes_rate_limit_ns("get_bytes_rate_limit_ns");
bvar::Adder<int64_t> get_bytes_rate_limit_exceed_req_num("get_bytes_rate_limit_exceed_req_num");
bvar::Adder<int64_t> put_bytes_rate_limit_ns("put_bytes_rate_limit_ns");
bvar::Adder<int64_t> put_bytes_rate_limit_exceed_req_num("put_bytes_rate_limit_exceed_req_num");

static std::atomic<int64_t> last_s3_get_token_bucket_tokens {0};
static std::atomic<int64_t> last_s3_get_token_limit {0};
static std::atomic<int64_t> last_s3_get_token_per_second {0};
static std::atomic<int64_t> last_s3_put_token_per_second {0};
static std::atomic<int64_t> last_s3_put_token_bucket_tokens {0};
static std::atomic<int64_t> last_s3_put_token_limit {0};
static std::atomic<int64_t> last_s3_get_bytes_per_second {0};
static std::atomic<int64_t> last_s3_get_bytes_bucket_tokens {0};
static std::atomic<int64_t> last_s3_put_bytes_per_second {0};
static std::atomic<int64_t> last_s3_put_bytes_bucket_tokens {0};

static std::atomic<bool> updating_get_limiter {false};
static std::atomic<bool> updating_put_limiter {false};
static std::atomic<bool> updating_get_bytes_limiter {false};
static std::atomic<bool> updating_put_bytes_limiter {false};

S3RateLimiterHolder* S3ClientFactory::rate_limiter(S3RateLimitType type) {
CHECK(type == S3RateLimitType::GET || type == S3RateLimitType::PUT) << to_string(type);
return _rate_limiters[static_cast<size_t>(type)].get();
}

S3RateLimiterHolder* S3ClientFactory::bytes_rate_limiter(S3RateLimitType type) {
CHECK(type == S3RateLimitType::GET || type == S3RateLimitType::PUT) << to_string(type);
return _bytes_rate_limiters[static_cast<size_t>(type)].get();
}

S3RateLimiterConfig get_effective_s3_rate_limiter_config(S3RateLimitType type) {
const int64_t cores = std::max<int64_t>(1, CpuInfo::num_cores());
S3RateLimiterConfig result;

switch (type) {
case S3RateLimitType::GET:
result.token_limit = config::s3_get_token_limit;
if (config::s3_get_token_per_second_per_core > 0) {
result.token_per_second = compute_s3_token_per_second_from_core(
config::s3_get_token_per_second_per_core, cores);
result.bucket_tokens = result.token_per_second;
} else {
result.token_per_second = config::s3_get_token_per_second;
result.bucket_tokens = config::s3_get_bucket_tokens;
}
return result;
case S3RateLimitType::PUT:
result.token_limit = config::s3_put_token_limit;
if (config::s3_put_token_per_second_per_core > 0) {
result.token_per_second = compute_s3_token_per_second_from_core(
config::s3_put_token_per_second_per_core, cores);
result.bucket_tokens = result.token_per_second;
} else {
result.token_per_second = config::s3_put_token_per_second;
result.bucket_tokens = config::s3_put_bucket_tokens;
}
return result;
default:
CHECK(false) << "unknown S3 rate limiter type: " << to_string(type);
return result;
}
}

S3RateLimiterConfig get_effective_s3_bytes_rate_limiter_config(S3RateLimitType type) {
const int64_t cores = std::max<int64_t>(1, CpuInfo::num_cores());
int64_t per_core = 0;

switch (type) {
case S3RateLimitType::GET:
per_core = config::s3_get_bytes_per_second_per_core;
break;
case S3RateLimitType::PUT:
per_core = config::s3_put_bytes_per_second_per_core;
break;
default:
CHECK(false) << "unknown S3 rate limiter type: " << to_string(type);
}

S3RateLimiterConfig result;
if (per_core > 0) {
result.token_per_second = compute_s3_bytes_per_second_from_core(per_core, cores);
result.bucket_tokens = result.token_per_second;
}
return result;
}

template <S3RateLimitType LimiterType>
void update_rate_limiter_if_changed(int64_t current_tps, int64_t current_bucket,
int64_t current_limit, std::atomic<int64_t>& last_tps,
Expand All @@ -201,6 +294,42 @@ void update_rate_limiter_if_changed(int64_t current_tps, int64_t current_bucket,
last_tps.store(current_tps, std::memory_order_release);
last_bucket.store(current_bucket, std::memory_order_release);
last_limit.store(current_limit, std::memory_order_release);
LOG(INFO) << "Reset S3 " << limiter_name
<< " rate limiter, token_per_second=" << current_tps
<< ", bucket_tokens=" << current_bucket
<< ", token_limit=" << current_limit;
} else {
LOG(WARNING) << "Failed to reset S3 " << limiter_name
<< " rate limiter, error code: " << ret;
}
}

updating_flag.store(false, std::memory_order_release);
}
}

void update_bytes_rate_limiter_if_changed(S3RateLimiterHolder* limiter,
int64_t current_bytes_per_second, int64_t current_bucket,
std::atomic<int64_t>& last_bytes_per_second,
std::atomic<int64_t>& last_bucket,
std::atomic<bool>& updating_flag,
const char* limiter_name) {
if (last_bytes_per_second.load(std::memory_order_relaxed) != current_bytes_per_second ||
last_bucket.load(std::memory_order_relaxed) != current_bucket) {
bool expected = false;
if (!updating_flag.compare_exchange_strong(expected, true, std::memory_order_acq_rel)) {
return;
}
if (last_bytes_per_second.load(std::memory_order_acquire) != current_bytes_per_second ||
last_bucket.load(std::memory_order_acquire) != current_bucket) {
int ret = limiter->reset(current_bytes_per_second, current_bucket, 0);

if (ret == 0) {
last_bytes_per_second.store(current_bytes_per_second, std::memory_order_release);
last_bucket.store(current_bucket, std::memory_order_release);
LOG(INFO) << "Reset S3 " << limiter_name
<< " rate limiter, bytes_per_second=" << current_bytes_per_second
<< ", bucket_bytes=" << current_bucket;
} else {
LOG(WARNING) << "Failed to reset S3 " << limiter_name
<< " rate limiter, error code: " << ret;
Expand All @@ -212,15 +341,31 @@ void update_rate_limiter_if_changed(int64_t current_tps, int64_t current_bucket,
}

void check_s3_rate_limiter_config_changed() {
auto get_config = get_effective_s3_rate_limiter_config(S3RateLimitType::GET);
update_rate_limiter_if_changed<S3RateLimitType::GET>(
config::s3_get_token_per_second, config::s3_get_bucket_tokens,
config::s3_get_token_limit, last_s3_get_token_per_second,
last_s3_get_token_bucket_tokens, last_s3_get_token_limit, updating_get_limiter, "GET");
get_config.token_per_second, get_config.bucket_tokens, get_config.token_limit,
last_s3_get_token_per_second, last_s3_get_token_bucket_tokens, last_s3_get_token_limit,
updating_get_limiter, "GET");

auto put_config = get_effective_s3_rate_limiter_config(S3RateLimitType::PUT);
update_rate_limiter_if_changed<S3RateLimitType::PUT>(
config::s3_put_token_per_second, config::s3_put_bucket_tokens,
config::s3_put_token_limit, last_s3_put_token_per_second,
last_s3_put_token_bucket_tokens, last_s3_put_token_limit, updating_put_limiter, "PUT");
put_config.token_per_second, put_config.bucket_tokens, put_config.token_limit,
last_s3_put_token_per_second, last_s3_put_token_bucket_tokens, last_s3_put_token_limit,
updating_put_limiter, "PUT");

auto get_bytes_config = get_effective_s3_bytes_rate_limiter_config(S3RateLimitType::GET);
update_bytes_rate_limiter_if_changed(
S3ClientFactory::instance().bytes_rate_limiter(S3RateLimitType::GET),
get_bytes_config.token_per_second, get_bytes_config.bucket_tokens,
last_s3_get_bytes_per_second, last_s3_get_bytes_bucket_tokens,
updating_get_bytes_limiter, "GET bytes");

auto put_bytes_config = get_effective_s3_bytes_rate_limiter_config(S3RateLimitType::PUT);
update_bytes_rate_limiter_if_changed(
S3ClientFactory::instance().bytes_rate_limiter(S3RateLimitType::PUT),
put_bytes_config.token_per_second, put_bytes_config.bucket_tokens,
last_s3_put_bytes_per_second, last_s3_put_bytes_bucket_tokens,
updating_put_bytes_limiter, "PUT bytes");
}

int reset_s3_rate_limiter(S3RateLimitType type, size_t max_speed, size_t max_burst, size_t limit) {
Expand All @@ -239,15 +384,30 @@ S3ClientFactory::S3ClientFactory() {
};
Aws::InitAPI(_aws_options);
_ca_cert_file_path = get_valid_ca_cert_path(doris::split(config::ca_cert_file_paths, ";"));
auto get_init_config = get_effective_s3_rate_limiter_config(S3RateLimitType::GET);
auto put_init_config = get_effective_s3_rate_limiter_config(S3RateLimitType::PUT);
_rate_limiters = {
std::make_unique<S3RateLimiterHolder>(
config::s3_get_token_per_second, config::s3_get_bucket_tokens,
config::s3_get_token_limit,
get_init_config.token_per_second, get_init_config.bucket_tokens,
get_init_config.token_limit,
metric_func_factory(get_rate_limit_ns, get_rate_limit_exceed_req_num)),
std::make_unique<S3RateLimiterHolder>(
config::s3_put_token_per_second, config::s3_put_bucket_tokens,
config::s3_put_token_limit,
put_init_config.token_per_second, put_init_config.bucket_tokens,
put_init_config.token_limit,
metric_func_factory(put_rate_limit_ns, put_rate_limit_exceed_req_num))};
auto get_bytes_init_config = get_effective_s3_bytes_rate_limiter_config(S3RateLimitType::GET);
auto put_bytes_init_config = get_effective_s3_bytes_rate_limiter_config(S3RateLimitType::PUT);
_bytes_rate_limiters = {
std::make_unique<S3RateLimiterHolder>(
get_bytes_init_config.token_per_second, get_bytes_init_config.bucket_tokens,
get_bytes_init_config.token_limit,
metric_func_factory(get_bytes_rate_limit_ns,
get_bytes_rate_limit_exceed_req_num)),
std::make_unique<S3RateLimiterHolder>(
put_bytes_init_config.token_per_second, put_bytes_init_config.bucket_tokens,
put_bytes_init_config.token_limit,
metric_func_factory(put_bytes_rate_limit_ns,
put_bytes_rate_limit_exceed_req_num))};

#ifdef USE_AZURE
auto azureLogLevel =
Expand Down
Loading
Loading