diff --git a/be/src/common/config.cpp b/be/src/common/config.cpp index 9ec587bd100ec5..68465e32d885bf 100644 --- a/be/src/common/config.cpp +++ b/be/src/common/config.cpp @@ -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 diff --git a/be/src/common/config.h b/be/src/common/config.h index 012f09a735c0cd..413caefe388e8e 100644 --- a/be/src/common/config.h +++ b/be/src/common/config.h @@ -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 diff --git a/be/src/io/fs/azure_obj_storage_client.cpp b/be/src/io/fs/azure_obj_storage_client.cpp index 85b81c1deb23df..94f3cadba7e492 100644 --- a/be/src/io/fs/azure_obj_storage_client.cpp +++ b/be/src/io/fs/azure_obj_storage_client.cpp @@ -70,25 +70,21 @@ auto base64_encode_part_num(int part_num) { } template -auto s3_rate_limit(doris::S3RateLimitType op, Func callback) -> decltype(callback()) { - if (!doris::config::enable_s3_rate_limiter) { - return callback(); - } - auto sleep_duration = doris::S3ClientFactory::instance().rate_limiter(op)->add(1); - if (sleep_duration < 0) { +auto s3_rate_limit(doris::S3RateLimitType op, size_t bytes, Func callback) -> decltype(callback()) { + if (!doris::acquire_s3_rate_limit(op, bytes)) { throw std::runtime_error("Azure exceeds request limit"); } return callback(); } template -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 -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)); } constexpr char SAS_TOKEN_URL_TEMPLATE[] = "{}/{}/{}{}"; @@ -225,11 +221,13 @@ ObjectStorageResponse AzureObjStorageClient::put_object(const ObjectStoragePathO auto client = _client->GetBlockBlobClient(opts.key); return do_azure_client_call( [&]() { - s3_put_rate_limit([&]() { - SCOPED_BVAR_LATENCY(s3_bvar::s3_put_latency); - client.UploadFrom(reinterpret_cast(stream.data()), - stream.size()); - }); + s3_put_rate_limit( + [&]() { + SCOPED_BVAR_LATENCY(s3_bvar::s3_put_latency); + client.UploadFrom(reinterpret_cast(stream.data()), + stream.size()); + }, + stream.size()); }, opts, _tls_debug_context); } @@ -243,10 +241,12 @@ ObjectStorageUploadResponse AzureObjStorageClient::upload_part(const ObjectStora Azure::Core::IO::MemoryBodyStream memory_body( reinterpret_cast(stream.data()), stream.size()); // The blockId must be base64 encoded - s3_put_rate_limit([&]() { - SCOPED_BVAR_LATENCY(s3_bvar::s3_multi_part_upload_latency); - client.StageBlock(base64_encode_part_num(part_num), memory_body); - }); + s3_put_rate_limit( + [&]() { + SCOPED_BVAR_LATENCY(s3_bvar::s3_multi_part_upload_latency); + client.StageBlock(base64_encode_part_num(part_num), memory_body); + }, + stream.size()); }, opts, _tls_debug_context); return ObjectStorageUploadResponse { @@ -305,11 +305,13 @@ ObjectStorageResponse AzureObjStorageClient::get_object(const ObjectStoragePathO DownloadBlobToOptions download_opts; Azure::Core::Http::HttpRange range {static_cast(offset), bytes_read}; download_opts.Range = range; - auto resp = s3_get_rate_limit([&]() { - SCOPED_BVAR_LATENCY(s3_bvar::s3_get_latency); - return client.DownloadTo(reinterpret_cast(buffer), bytes_read, - download_opts); - }); + auto resp = s3_get_rate_limit( + [&]() { + SCOPED_BVAR_LATENCY(s3_bvar::s3_get_latency); + return client.DownloadTo(reinterpret_cast(buffer), bytes_read, + download_opts); + }, + bytes_read); *size_return = resp.Value.ContentRange.Length.Value(); }, opts, _tls_debug_context); diff --git a/be/src/io/fs/s3_obj_storage_client.cpp b/be/src/io/fs/s3_obj_storage_client.cpp index 1b8bf7b473c076..3c6798cb8d5f34 100644 --- a/be/src/io/fs/s3_obj_storage_client.cpp +++ b/be/src/io/fs/s3_obj_storage_client.cpp @@ -77,26 +77,22 @@ inline ::Aws::Client::AWSError<::Aws::S3::S3Errors> s3_error_factory() { } template -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(); - } - auto sleep_duration = doris::S3ClientFactory::instance().rate_limiter(op)->add(1); - if (sleep_duration < 0) { + if (!doris::acquire_s3_rate_limit(op, bytes)) { return T(s3_error_factory()); } return callback(); } template -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 -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 @@ -167,7 +163,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(); @@ -211,7 +207,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(); @@ -322,7 +318,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))), @@ -502,4 +498,4 @@ std::string S3ObjStorageClient::generate_presigned_url(const ObjectStoragePathOp expiration_secs); } -} // namespace doris::io \ No newline at end of file +} // namespace doris::io diff --git a/be/src/util/s3_util.cpp b/be/src/util/s3_util.cpp index fc4c81d4bd96f2..5b3b16d84560b0 100644 --- a/be/src/util/s3_util.cpp +++ b/be/src/util/s3_util.cpp @@ -33,7 +33,9 @@ #include #include +#include +#include "util/cpu_info.h" #include "util/string_util.h" #ifdef USE_AZURE @@ -156,12 +158,36 @@ 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::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::max(); + if (per_core > max_bytes_per_second / cores) { + return max_bytes_per_second; + } + return per_core * cores; +} } // namespace bvar::Adder get_rate_limit_ns("get_rate_limit_ns"); bvar::Adder get_rate_limit_exceed_req_num("get_rate_limit_exceed_req_num"); bvar::Adder put_rate_limit_ns("put_rate_limit_ns"); bvar::Adder put_rate_limit_exceed_req_num("put_rate_limit_exceed_req_num"); +bvar::Adder get_bytes_rate_limit_ns("get_bytes_rate_limit_ns"); +bvar::Adder get_bytes_rate_limit_exceed_req_num("get_bytes_rate_limit_exceed_req_num"); +bvar::Adder put_bytes_rate_limit_ns("put_bytes_rate_limit_ns"); +bvar::Adder put_bytes_rate_limit_exceed_req_num("put_bytes_rate_limit_exceed_req_num"); static std::atomic last_s3_get_token_bucket_tokens {0}; static std::atomic last_s3_get_token_limit {0}; @@ -169,15 +195,94 @@ static std::atomic last_s3_get_token_per_second {0}; static std::atomic last_s3_put_token_per_second {0}; static std::atomic last_s3_put_token_bucket_tokens {0}; static std::atomic last_s3_put_token_limit {0}; +static std::atomic last_s3_get_bytes_per_second {0}; +static std::atomic last_s3_get_bytes_bucket_tokens {0}; +static std::atomic last_s3_put_bytes_per_second {0}; +static std::atomic last_s3_put_bytes_bucket_tokens {0}; static std::atomic updating_get_limiter {false}; static std::atomic updating_put_limiter {false}; +static std::atomic updating_get_bytes_limiter {false}; +static std::atomic 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(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(type)].get(); +} + +bool acquire_s3_rate_limit(S3RateLimitType type, size_t bytes) { + if (!config::enable_s3_rate_limiter) { + return true; + } + + auto& factory = S3ClientFactory::instance(); + if (factory.rate_limiter(type)->add(1) < 0) { + return false; + } + return bytes == 0 || factory.bytes_rate_limiter(type)->add(bytes) >= 0; +} + +S3RateLimiterConfig get_effective_s3_rate_limiter_config(S3RateLimitType type) { + const int64_t cores = std::max(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(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 void update_rate_limiter_if_changed(int64_t current_tps, int64_t current_bucket, int64_t current_limit, std::atomic& last_tps, @@ -201,6 +306,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& last_bytes_per_second, + std::atomic& last_bucket, + std::atomic& 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; @@ -212,15 +353,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( - 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( - 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) { @@ -239,15 +396,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( - 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( - 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( + 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( + 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 = diff --git a/be/src/util/s3_util.h b/be/src/util/s3_util.h index 5546db857c91b1..b4402ffe2c385e 100644 --- a/be/src/util/s3_util.h +++ b/be/src/util/s3_util.h @@ -64,10 +64,21 @@ extern bvar::LatencyRecorder s3_copy_object_latency; std::string hide_access_key(const std::string& ak); int reset_s3_rate_limiter(S3RateLimitType type, size_t max_speed, size_t max_burst, size_t limit); -// Rebuild the S3 GET/PUT rate limiters if the related configs have changed. +// Rebuild the S3 GET/PUT request and bytes rate limiters if related configs have changed. // Safe to call periodically; it is a no-op when nothing changed. void check_s3_rate_limiter_config_changed(); +bool acquire_s3_rate_limit(S3RateLimitType type, size_t bytes); + +struct S3RateLimiterConfig { + int64_t token_per_second = 0; + int64_t bucket_tokens = 0; + int64_t token_limit = 0; +}; + +S3RateLimiterConfig get_effective_s3_rate_limiter_config(S3RateLimitType type); +S3RateLimiterConfig get_effective_s3_bytes_rate_limiter_config(S3RateLimitType type); + class S3URI; struct S3ClientConf { std::string endpoint; @@ -158,6 +169,7 @@ class S3ClientFactory { } S3RateLimiterHolder* rate_limiter(S3RateLimitType type); + S3RateLimiterHolder* bytes_rate_limiter(S3RateLimitType type); std::shared_ptr get_aws_credentials_provider( const S3ClientConf& s3_conf); @@ -186,6 +198,7 @@ class S3ClientFactory { std::unordered_map> _cache; std::string _ca_cert_file_path; std::array, 2> _rate_limiters; + std::array, 2> _bytes_rate_limiters; #ifdef BE_TEST std::function(const S3ClientConf&)> _test_client_creator; #endif diff --git a/be/test/util/s3_util_test.cpp b/be/test/util/s3_util_test.cpp index a87d0e9a2d2980..40ddd98edfa5ba 100644 --- a/be/test/util/s3_util_test.cpp +++ b/be/test/util/s3_util_test.cpp @@ -19,14 +19,65 @@ #include +#include #include #include "common/config.h" #include "gtest/gtest_pred_impl.h" +#include "util/cpu_info.h" #include "util/s3_uri.h" namespace doris { +struct S3RateLimiterConfigGuard { + bool enable_rate_limiter = config::enable_s3_rate_limiter; + int64_t get_tps = config::s3_get_token_per_second; + int64_t get_bucket = config::s3_get_bucket_tokens; + int64_t get_limit = config::s3_get_token_limit; + int64_t put_tps = config::s3_put_token_per_second; + int64_t put_bucket = config::s3_put_bucket_tokens; + int64_t put_limit = config::s3_put_token_limit; + int64_t get_per_core = config::s3_get_token_per_second_per_core; + int64_t put_per_core = config::s3_put_token_per_second_per_core; + int64_t max_tps = config::s3_token_per_second_max; + int64_t get_bytes_per_core = config::s3_get_bytes_per_second_per_core; + int64_t put_bytes_per_core = config::s3_put_bytes_per_second_per_core; + int64_t max_bytes_per_second = config::s3_bytes_per_second_max; + + ~S3RateLimiterConfigGuard() { + config::enable_s3_rate_limiter = enable_rate_limiter; + config::s3_get_token_per_second = get_tps; + config::s3_get_bucket_tokens = get_bucket; + config::s3_get_token_limit = get_limit; + config::s3_put_token_per_second = put_tps; + config::s3_put_bucket_tokens = put_bucket; + config::s3_put_token_limit = put_limit; + config::s3_get_token_per_second_per_core = get_per_core; + config::s3_put_token_per_second_per_core = put_per_core; + config::s3_token_per_second_max = max_tps; + config::s3_get_bytes_per_second_per_core = get_bytes_per_core; + config::s3_put_bytes_per_second_per_core = put_bytes_per_core; + config::s3_bytes_per_second_max = max_bytes_per_second; + check_s3_rate_limiter_config_changed(); + + auto get_config = get_effective_s3_rate_limiter_config(S3RateLimitType::GET); + auto put_config = get_effective_s3_rate_limiter_config(S3RateLimitType::PUT); + auto get_bytes_config = get_effective_s3_bytes_rate_limiter_config(S3RateLimitType::GET); + auto put_bytes_config = get_effective_s3_bytes_rate_limiter_config(S3RateLimitType::PUT); + auto& factory = S3ClientFactory::instance(); + factory.rate_limiter(S3RateLimitType::GET) + ->reset(get_config.token_per_second, get_config.bucket_tokens, + get_config.token_limit); + factory.rate_limiter(S3RateLimitType::PUT) + ->reset(put_config.token_per_second, put_config.bucket_tokens, + put_config.token_limit); + factory.bytes_rate_limiter(S3RateLimitType::GET) + ->reset(get_bytes_config.token_per_second, get_bytes_config.bucket_tokens, 0); + factory.bytes_rate_limiter(S3RateLimitType::PUT) + ->reset(put_bytes_config.token_per_second, put_bytes_config.bucket_tokens, 0); + } +}; + class S3UTILTest : public testing::Test { public: S3UTILTest() = default; @@ -82,14 +133,10 @@ TEST_F(S3UTILTest, hide_access_key_typical_aws_key) { // thread relies on to apply dynamically modified s3_get_* rate limiter configs without // having to (re)create an S3 client. TEST_F(S3UTILTest, check_s3_rate_limiter_config_changed_rebuilds_limiter) { + S3RateLimiterConfigGuard guard; auto* get_limiter = S3ClientFactory::instance().rate_limiter(S3RateLimitType::GET); ASSERT_NE(get_limiter, nullptr); - // Save originals so other tests are not affected. - const int64_t orig_tps = config::s3_get_token_per_second; - const int64_t orig_bucket = config::s3_get_bucket_tokens; - const int64_t orig_limit = config::s3_get_token_limit; - // Establish a known baseline (no count limit, no throttling). config::s3_get_token_per_second = 1000000000; config::s3_get_bucket_tokens = 1000000000; @@ -114,11 +161,252 @@ TEST_F(S3UTILTest, check_s3_rate_limiter_config_changed_rebuilds_limiter) { check_s3_rate_limiter_config_changed(); EXPECT_GE(get_limiter->add(1), 0); - // Restore original configs and apply them back to the limiter. - config::s3_get_token_per_second = orig_tps; - config::s3_get_bucket_tokens = orig_bucket; - config::s3_get_token_limit = orig_limit; + // guard restores original configs and re-applies them on destruction. +} + +TEST_F(S3UTILTest, acquire_s3_rate_limit_charges_request_and_bytes_limiters) { + S3RateLimiterConfigGuard guard; + config::enable_s3_rate_limiter = true; + + auto* request_limiter = S3ClientFactory::instance().rate_limiter(S3RateLimitType::GET); + auto* bytes_limiter = S3ClientFactory::instance().bytes_rate_limiter(S3RateLimitType::GET); + ASSERT_NE(request_limiter, nullptr); + ASSERT_NE(bytes_limiter, nullptr); + + request_limiter->reset(1000000000, 1000000000, 2); + bytes_limiter->reset(1000000000, 1000000000, 5); + + EXPECT_TRUE(acquire_s3_rate_limit(S3RateLimitType::GET, 3)); + EXPECT_FALSE(acquire_s3_rate_limit(S3RateLimitType::GET, 3)); +} + +TEST_F(S3UTILTest, effective_s3_rate_limiter_config_uses_absolute_get_config) { + S3RateLimiterConfigGuard guard; + config::s3_get_token_per_second = 123; + config::s3_get_bucket_tokens = 456; + config::s3_get_token_limit = 789; + config::s3_get_token_per_second_per_core = 0; + config::s3_token_per_second_max = 9; + + auto effective = get_effective_s3_rate_limiter_config(S3RateLimitType::GET); + EXPECT_EQ(123, effective.token_per_second); + EXPECT_EQ(456, effective.bucket_tokens); + EXPECT_EQ(789, effective.token_limit); +} + +TEST_F(S3UTILTest, effective_s3_rate_limiter_config_uses_absolute_put_config) { + S3RateLimiterConfigGuard guard; + config::s3_put_token_per_second = 321; + config::s3_put_bucket_tokens = 654; + config::s3_put_token_limit = 987; + config::s3_put_token_per_second_per_core = 0; + config::s3_token_per_second_max = 9; + + auto effective = get_effective_s3_rate_limiter_config(S3RateLimitType::PUT); + EXPECT_EQ(321, effective.token_per_second); + EXPECT_EQ(654, effective.bucket_tokens); + EXPECT_EQ(987, effective.token_limit); +} + +TEST_F(S3UTILTest, effective_s3_rate_limiter_config_uses_per_core_get_config_without_cap) { + S3RateLimiterConfigGuard guard; + const int64_t cores = std::max(1, CpuInfo::num_cores()); + config::s3_get_token_per_second = 123; + config::s3_get_bucket_tokens = 456; + config::s3_get_token_limit = 789; + config::s3_get_token_per_second_per_core = 8; + config::s3_token_per_second_max = 0; + + auto effective = get_effective_s3_rate_limiter_config(S3RateLimitType::GET); + EXPECT_EQ(8 * cores, effective.token_per_second); + EXPECT_EQ(8 * cores, effective.bucket_tokens); + EXPECT_EQ(789, effective.token_limit); +} + +TEST_F(S3UTILTest, effective_s3_rate_limiter_config_caps_per_core_get_config) { + S3RateLimiterConfigGuard guard; + config::s3_get_token_limit = 789; + config::s3_get_token_per_second_per_core = 1000000; + config::s3_token_per_second_max = 256; + + auto effective = get_effective_s3_rate_limiter_config(S3RateLimitType::GET); + EXPECT_EQ(256, effective.token_per_second); + EXPECT_EQ(256, effective.bucket_tokens); + EXPECT_EQ(789, effective.token_limit); +} + +TEST_F(S3UTILTest, effective_s3_rate_limiter_config_caps_per_core_put_config) { + S3RateLimiterConfigGuard guard; + config::s3_put_token_limit = 987; + config::s3_put_token_per_second_per_core = 1000000; + config::s3_token_per_second_max = 128; + + auto effective = get_effective_s3_rate_limiter_config(S3RateLimitType::PUT); + EXPECT_EQ(128, effective.token_per_second); + EXPECT_EQ(128, effective.bucket_tokens); + EXPECT_EQ(987, effective.token_limit); +} + +TEST_F(S3UTILTest, effective_s3_rate_limiter_config_caps_overflowing_per_core_get_config) { + S3RateLimiterConfigGuard guard; + config::s3_get_token_limit = 789; + config::s3_get_token_per_second_per_core = std::numeric_limits::max(); + config::s3_token_per_second_max = 1024; + + auto effective = get_effective_s3_rate_limiter_config(S3RateLimitType::GET); + EXPECT_EQ(1024, effective.token_per_second); + EXPECT_EQ(1024, effective.bucket_tokens); + EXPECT_EQ(789, effective.token_limit); +} + +TEST_F(S3UTILTest, check_s3_rate_limiter_config_changed_rebuilds_get_limiter_for_per_core_config) { + S3RateLimiterConfigGuard guard; + auto* get_limiter = S3ClientFactory::instance().rate_limiter(S3RateLimitType::GET); + ASSERT_NE(get_limiter, nullptr); + + const int64_t cores = std::max(1, CpuInfo::num_cores()); + config::s3_get_token_limit = 0; + config::s3_get_token_per_second_per_core = 7; + config::s3_token_per_second_max = 0; check_s3_rate_limiter_config_changed(); + + EXPECT_EQ(7 * cores, get_limiter->get_max_speed()); + EXPECT_EQ(7 * cores, get_limiter->get_max_burst()); + EXPECT_EQ(0, get_limiter->get_limit()); +} + +TEST_F(S3UTILTest, check_s3_rate_limiter_config_changed_rebuilds_get_limiter_when_max_changes) { + S3RateLimiterConfigGuard guard; + auto* get_limiter = S3ClientFactory::instance().rate_limiter(S3RateLimitType::GET); + ASSERT_NE(get_limiter, nullptr); + + config::s3_get_token_limit = 0; + config::s3_get_token_per_second_per_core = 1000000; + config::s3_token_per_second_max = 512; + check_s3_rate_limiter_config_changed(); + + EXPECT_EQ(512, get_limiter->get_max_speed()); + EXPECT_EQ(512, get_limiter->get_max_burst()); + EXPECT_EQ(0, get_limiter->get_limit()); + + config::s3_token_per_second_max = 256; + check_s3_rate_limiter_config_changed(); + + EXPECT_EQ(256, get_limiter->get_max_speed()); + EXPECT_EQ(256, get_limiter->get_max_burst()); + EXPECT_EQ(0, get_limiter->get_limit()); +} + +TEST_F(S3UTILTest, check_s3_rate_limiter_config_changed_rebuilds_put_limiter_when_max_changes) { + S3RateLimiterConfigGuard guard; + auto* put_limiter = S3ClientFactory::instance().rate_limiter(S3RateLimitType::PUT); + ASSERT_NE(put_limiter, nullptr); + + config::s3_put_token_limit = 0; + config::s3_put_token_per_second_per_core = 1000000; + config::s3_token_per_second_max = 512; + check_s3_rate_limiter_config_changed(); + + EXPECT_EQ(512, put_limiter->get_max_speed()); + EXPECT_EQ(512, put_limiter->get_max_burst()); + EXPECT_EQ(0, put_limiter->get_limit()); + + config::s3_token_per_second_max = 256; + check_s3_rate_limiter_config_changed(); + + EXPECT_EQ(256, put_limiter->get_max_speed()); + EXPECT_EQ(256, put_limiter->get_max_burst()); + EXPECT_EQ(0, put_limiter->get_limit()); +} + +TEST_F(S3UTILTest, check_s3_rate_limiter_config_changed_rebuilds_put_limiter_for_per_core_config) { + S3RateLimiterConfigGuard guard; + auto* put_limiter = S3ClientFactory::instance().rate_limiter(S3RateLimitType::PUT); + ASSERT_NE(put_limiter, nullptr); + + const int64_t cores = std::max(1, CpuInfo::num_cores()); + config::s3_put_token_limit = 0; + config::s3_put_token_per_second_per_core = 5; + config::s3_token_per_second_max = 0; + check_s3_rate_limiter_config_changed(); + + EXPECT_EQ(5 * cores, put_limiter->get_max_speed()); + EXPECT_EQ(5 * cores, put_limiter->get_max_burst()); + EXPECT_EQ(0, put_limiter->get_limit()); +} + +TEST_F(S3UTILTest, effective_s3_bytes_rate_limiter_config_is_disabled_when_per_core_is_zero) { + S3RateLimiterConfigGuard guard; + config::s3_get_bytes_per_second_per_core = 0; + config::s3_bytes_per_second_max = 9; + + auto effective = get_effective_s3_bytes_rate_limiter_config(S3RateLimitType::GET); + EXPECT_EQ(0, effective.token_per_second); + EXPECT_EQ(0, effective.bucket_tokens); + EXPECT_EQ(0, effective.token_limit); +} + +TEST_F(S3UTILTest, effective_s3_bytes_rate_limiter_config_uses_per_core_get_config) { + S3RateLimiterConfigGuard guard; + const int64_t cores = std::max(1, CpuInfo::num_cores()); + config::s3_get_bytes_per_second_per_core = 1024; + config::s3_bytes_per_second_max = 0; + + auto effective = get_effective_s3_bytes_rate_limiter_config(S3RateLimitType::GET); + EXPECT_EQ(1024 * cores, effective.token_per_second); + EXPECT_EQ(1024 * cores, effective.bucket_tokens); + EXPECT_EQ(0, effective.token_limit); +} + +TEST_F(S3UTILTest, effective_s3_bytes_rate_limiter_config_caps_per_core_put_config) { + S3RateLimiterConfigGuard guard; + config::s3_put_bytes_per_second_per_core = 1000000; + config::s3_bytes_per_second_max = 4096; + + auto effective = get_effective_s3_bytes_rate_limiter_config(S3RateLimitType::PUT); + EXPECT_EQ(4096, effective.token_per_second); + EXPECT_EQ(4096, effective.bucket_tokens); + EXPECT_EQ(0, effective.token_limit); +} + +TEST_F(S3UTILTest, check_s3_rate_limiter_config_changed_rebuilds_get_bytes_limiter) { + S3RateLimiterConfigGuard guard; + auto* get_limiter = S3ClientFactory::instance().bytes_rate_limiter(S3RateLimitType::GET); + ASSERT_NE(get_limiter, nullptr); + + config::s3_get_bytes_per_second_per_core = 1000000; + config::s3_bytes_per_second_max = 8192; + check_s3_rate_limiter_config_changed(); + + EXPECT_EQ(8192, get_limiter->get_max_speed()); + EXPECT_EQ(8192, get_limiter->get_max_burst()); + EXPECT_EQ(0, get_limiter->get_limit()); + + config::s3_bytes_per_second_max = 4096; + check_s3_rate_limiter_config_changed(); + + EXPECT_EQ(4096, get_limiter->get_max_speed()); + EXPECT_EQ(4096, get_limiter->get_max_burst()); + EXPECT_EQ(0, get_limiter->get_limit()); +} + +TEST_F(S3UTILTest, bytes_limiter_config_check_does_not_reset_qps_limiter) { + S3RateLimiterConfigGuard guard; + auto* get_limiter = S3ClientFactory::instance().rate_limiter(S3RateLimitType::GET); + ASSERT_NE(get_limiter, nullptr); + + config::s3_get_token_per_second = 1000000000; + config::s3_get_bucket_tokens = 1000000000; + config::s3_get_token_limit = 3; + config::s3_get_token_per_second_per_core = 0; + config::s3_get_bytes_per_second_per_core = 0; + check_s3_rate_limiter_config_changed(); + + EXPECT_GE(get_limiter->add(1), 0); + check_s3_rate_limiter_config_changed(); + EXPECT_GE(get_limiter->add(1), 0); + EXPECT_GE(get_limiter->add(1), 0); + EXPECT_LT(get_limiter->add(1), 0); } } // end namespace doris diff --git a/common/cpp/token_bucket_rate_limiter.h b/common/cpp/token_bucket_rate_limiter.h index dbeca3b897db03..69671a309db8f6 100644 --- a/common/cpp/token_bucket_rate_limiter.h +++ b/common/cpp/token_bucket_rate_limiter.h @@ -61,7 +61,7 @@ class TokenBucketRateLimiter { private: std::pair _update_remain_token(long now, size_t amount); size_t _count {0}; - const size_t _max_speed {0}; // in tokens per second. which indicates the QPS + const size_t _max_speed {0}; // Maximum number of tokens added per second. const size_t _max_burst {0}; // in tokens. which indicates the token bucket size const uint64_t _limit {0}; // 0 - not limited. class SimpleSpinLock;