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

be/src/cloud/cloud_storage_engine.cpp

@@ -453,15 +453,9 @@ void CloudStorageEngine::_refresh_storage_vault_info_thread_callback() {
 while (!_stop_background_threads_latch.wait_for(
 std::chrono::seconds(config::refresh_s3_info_interval_s))) {
 sync_storage_vault();
-// The other place that rebuilds the S3 rate limiter is S3ClientFactory::create(), which
-// is not called when an existing vault's conf is unchanged. Trigger the check here as well
-// so that dynamically modified s3_{get,put}_* rate limiter configs take effect within
-// refresh_s3_info_interval_s even when no vault is created or its conf does not change.
-// Gate it behind enable_s3_rate_limiter so that clusters with rate limiting disabled
-// (e.g. HDFS-only vaults) do not force-initialize S3ClientFactory / the AWS SDK here.
-if (config::enable_s3_rate_limiter) {
-check_s3_rate_limiter_config_changed();
- }
+// Dynamically modified s3_{get,put}_* rate limiter configs and cgroup CPU quota
+// changes are picked up by the daemon's s3_rate_limiter_refresh_thread, which
+// runs in both cloud and non-cloud mode.
  }
 }
 

be/src/common/config.cpp

@@ -1550,6 +1550,34 @@ DEFINE_mInt64(s3_put_token_limit, "0");
 DEFINE_mInt64(s3_rate_limiter_log_interval, "1000");
 DEFINE_Validator(s3_rate_limiter_log_interval, [](int64_t config) -> bool { return config >= 0; });
 
+// CPU-aware S3 rate limiter. Effective GET/PUT requests per second =
+// requests_per_second_per_core * BE cpu cores, capped by the corresponding
+// requests_per_second_max. A negative value means unset: fall back to the legacy absolute
+// s3_{get,put}_token_* configs above. 0 disables request-rate limiting for that operation.
+DEFINE_mInt64(s3_get_requests_per_second_per_core, "-1");
+DEFINE_mInt64(s3_put_requests_per_second_per_core, "-1");
+// Hard caps for the CPU-derived GET/PUT QPS. A non-positive value means no cap.
+DEFINE_mInt64(s3_get_requests_per_second_max, "0");
+DEFINE_mInt64(s3_put_requests_per_second_max, "0");
+
+// CPU-aware S3 bandwidth limiter. Effective GET/PUT bytes/s = bytes_per_second_per_core *
+// BE cpu cores, capped by the corresponding bytes_per_second_max. A non-positive value disables
+// byte-rate limiting for that operation (there is no legacy fallback for bandwidth).
+// Note: the derived per-BE bytes/s should not be set below the single IO upper bound
+// per second (s3_write_buffer_size, 5MB by default). A single IO larger than 1 second
+// of quota only reserves 1 second worth of tokens; the excess bytes are not accounted
+// (reservation clamp in S3RateLimitGuard).
+DEFINE_mInt64(s3_get_bytes_per_second_per_core, "-1");
+DEFINE_mInt64(s3_put_bytes_per_second_per_core, "-1");
+// Hard caps for the CPU-derived GET/PUT bytes/s. A non-positive value means no cap.
+DEFINE_mInt64(s3_get_bytes_per_second_max, "0");
+DEFINE_mInt64(s3_put_bytes_per_second_max, "0");
+
+// Override the CPU cores used to derive the effective S3 rate limits. A non-positive value
+// means auto-detect from the cgroup cpu quota (fall back to physical cores); the control plane
+// can push a positive value via /api/update_config when resizing a serverless BE.
+DEFINE_mInt32(s3_rate_limiter_cpu_cores_override, "0");
+
 // The dir TrinoConnectorPluginLoader loads Trino's own plugins from, used verbatim. Keep the default
 // in sync with FE Config.trino_connector_plugin_dir: FE and BE load the same plugins and an operator
 // who leaves both untouched expects both to find them.

be/src/common/config.h

@@ -1621,6 +1621,22 @@ DECLARE_mInt64(s3_put_bucket_tokens);
 DECLARE_mInt64(s3_put_token_per_second);
 DECLARE_mInt64(s3_put_token_limit);
 DECLARE_mInt64(s3_rate_limiter_log_interval);
+
+// CPU-aware S3 rate limiter: GET/PUT QPS per CPU core. A negative value means unset and
+// falls back to the legacy absolute token configs above; 0 disables QPS limiting.
+DECLARE_mInt64(s3_get_requests_per_second_per_core);
+DECLARE_mInt64(s3_put_requests_per_second_per_core);
+// Hard caps for the CPU-derived GET/PUT QPS. A non-positive value means no cap.
+DECLARE_mInt64(s3_get_requests_per_second_max);
+DECLARE_mInt64(s3_put_requests_per_second_max);
+// GET/PUT bytes per second per CPU core. A non-positive value disables byte-rate limiting.
+DECLARE_mInt64(s3_get_bytes_per_second_per_core);
+DECLARE_mInt64(s3_put_bytes_per_second_per_core);
+// Hard caps for the CPU-derived GET/PUT bytes/s. A non-positive value means no cap.
+DECLARE_mInt64(s3_get_bytes_per_second_max);
+DECLARE_mInt64(s3_put_bytes_per_second_max);
+// Override for cores used to derive effective limits: a non-positive value means auto-detect.
+DECLARE_mInt32(s3_rate_limiter_cpu_cores_override);
 // max s3 client retry times
 DECLARE_mInt32(max_s3_client_retry);
 // When meet s3 429 error, the "get" request will

be/src/common/daemon.cpp

@@ -60,6 +60,7 @@
 #include "util/algorithm_util.h"
 #include "util/mem_info.h"
 #include "util/perf_counters.h"
+#include "util/s3_rate_limiter_manager.h"
 #include "util/time.h"
 
 namespace doris {
@@ -582,6 +583,18 @@ void Daemon::calculate_workload_group_metrics_thread() {
  }
 }
 
+void Daemon::s3_rate_limiter_refresh_thread() {
+ // Single trigger for dynamic rate limiter changes: picks up both mutable
+ // s3_{get,put}_* config updates and cgroup CPU quota changes (serverless BEs can
+ // be resized in place). refresh() is idempotent and compares against the buckets'
+ // own parameters, so quiet iterations are cheap no-ops.
+ while (!_stop_background_threads_latch.wait_for(std::chrono::seconds(10))) {
+ if (config::enable_s3_rate_limiter) {
+ S3RateLimiterManager::instance().refresh();
+ }
+ }
+}
+
 void Daemon::start() {
  Status st;
  st = Thread::create(
@@ -603,6 +616,10 @@ void Daemon::start() {
  [this]() { this->calculate_metrics_thread(); }, &_threads.emplace_back());
  CHECK(st.ok()) << st;
  }
+ st = Thread::create(
+ "Daemon", "s3_rate_limiter_refresh_thread",
+ [this]() { this->s3_rate_limiter_refresh_thread(); }, &_threads.emplace_back());
+ CHECK(st.ok()) << st;
  st = Thread::create(
  "Daemon", "je_reset_dirty_decay_thread",
  [this]() { this->je_reset_dirty_decay_thread(); }, &_threads.emplace_back());

be/src/common/daemon.h

@@ -46,6 +46,7 @@ class Daemon {
 voidreport_runtime_query_statistics_thread();
 voidbe_proc_monitor_thread();
 voidcalculate_workload_group_metrics_thread();
+voids3_rate_limiter_refresh_thread();
 
  CountDownLatch _stop_background_threads_latch;
  std::vector<std::shared_ptr<Thread>> _threads;

be/src/io/fs/azure_obj_storage_client.cpp

@@ -70,27 +70,8 @@ auto base64_encode_part_num(int part_num) {
  return Aws::Utils::HashingUtils::Base64Encode({buf, sizeof(buf)});
 }
 
-template <typename Func>
-auto s3_rate_limit(doris::S3RateLimitType op, Func callback) -> decltype(callback()) {
- if (!doris::config::enable_s3_rate_limiter) {
- return callback();
- }
- auto sleep_duration = doris::apply_s3_rate_limit(op);
- if (sleep_duration < 0) {
- throw std::runtime_error("Azure exceeds request limit");
- }
- return callback();
-}
-
-template <typename Func>
-auto s3_get_rate_limit(Func callback) -> decltype(callback()) {
- return s3_rate_limit(doris::S3RateLimitType::GET, 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));
-}
+// Rate limiting is applied by RateLimitedObjStorageClient, the decorator that
+// S3ClientFactory wraps around this client when the bucket is subject to limiting.
 
 constexpr char SAS_TOKEN_URL_TEMPLATE[] = "{}/{}/{}{}";
 constexpr char BlobNotFound[] = "BlobNotFound";
@@ -163,10 +144,8 @@ struct AzureBatchDeleter {
  }
  auto resp = do_azure_client_call(
  [&]() {
- s3_put_rate_limit([&]() {
- SCOPED_BVAR_LATENCY(s3_bvar::s3_delete_objects_latency);
- _client->SubmitBatch(_batch);
- });
+ SCOPED_BVAR_LATENCY(s3_bvar::s3_delete_objects_latency);
+ _client->SubmitBatch(_batch);
  },
  _opts, _tls_debug_context);
  if (resp.status.code != ErrorCode::OK) {
@@ -228,11 +207,8 @@ 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<const uint8_t*>(stream.data()),
- stream.size());
- });
+ SCOPED_BVAR_LATENCY(s3_bvar::s3_put_latency);
+ client.UploadFrom(reinterpret_cast<const uint8_t*>(stream.data()), stream.size());
  },
  opts, _tls_debug_context);
 }
@@ -246,10 +222,8 @@ ObjectStorageUploadResponse AzureObjStorageClient::upload_part(const ObjectStora
  Azure::Core::IO::MemoryBodyStream memory_body(
  reinterpret_cast<const uint8_t*>(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);
- });
+ SCOPED_BVAR_LATENCY(s3_bvar::s3_multi_part_upload_latency);
+ client.StageBlock(base64_encode_part_num(part_num), memory_body);
  },
  opts, _tls_debug_context);
  return ObjectStorageUploadResponse {
@@ -267,10 +241,8 @@ ObjectStorageResponse AzureObjStorageClient::complete_multipart_upload(
  [](const ObjectCompleteMultiPart& i) { return base64_encode_part_num(i.part_num); });
  return do_azure_client_call(
  [&]() {
- s3_put_rate_limit([&]() {
- SCOPED_BVAR_LATENCY(s3_bvar::s3_multi_part_upload_latency);
- client.CommitBlockList(string_block_ids);
- });
+ SCOPED_BVAR_LATENCY(s3_bvar::s3_multi_part_upload_latency);
+ client.CommitBlockList(string_block_ids);
  },
  opts, _tls_debug_context);
 }
@@ -279,10 +251,8 @@ ObjectStorageHeadResponse AzureObjStorageClient::head_object(const ObjectStorage
  Models::BlobProperties properties {};
  auto resp = do_azure_client_call(
  [&]() {
- properties = s3_get_rate_limit([&]() {
- SCOPED_BVAR_LATENCY(s3_bvar::s3_head_latency);
- return _client->GetBlockBlobClient(opts.key).GetProperties().Value;
- });
+ SCOPED_BVAR_LATENCY(s3_bvar::s3_head_latency);
+ properties = _client->GetBlockBlobClient(opts.key).GetProperties().Value;
  },
  opts, _tls_debug_context);
  if (resp.http_code == static_cast<int>(Azure::Core::Http::HttpStatusCode::NotFound)) {
@@ -308,11 +278,9 @@ ObjectStorageResponse AzureObjStorageClient::get_object(const ObjectStoragePathO
  DownloadBlobToOptions download_opts;
  Azure::Core::Http::HttpRange range {static_cast<int64_t>(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<uint8_t*>(buffer), bytes_read,
- download_opts);
- });
+ SCOPED_BVAR_LATENCY(s3_bvar::s3_get_latency);
+ auto resp = client.DownloadTo(reinterpret_cast<uint8_t*>(buffer), bytes_read,
+ download_opts);
  *size_return = resp.Value.ContentRange.Length.Value();
  },
  opts, _tls_debug_context);
@@ -330,17 +298,18 @@ ObjectStorageResponse AzureObjStorageClient::list_objects(const ObjectStoragePat
  [&]() {
  ListBlobsOptions list_opts;
  list_opts.Prefix = opts.prefix;
- auto resp = s3_get_rate_limit([&]() {
+ ListBlobsPagedResponse resp;
+ {
  SCOPED_BVAR_LATENCY(s3_bvar::s3_list_latency);
-return _client->ListBlobs(list_opts);
- });
+resp = _client->ListBlobs(list_opts);
+ }
  get_file_file(resp);
  while (resp.NextPageToken.HasValue()) {
  list_opts.ContinuationToken = resp.NextPageToken;
-resp = s3_get_rate_limit([&]() {
+ {
  SCOPED_BVAR_LATENCY(s3_bvar::s3_list_latency);
-return _client->ListBlobs(list_opts);
- });
+resp = _client->ListBlobs(list_opts);
+ }
  get_file_file(resp);
  }
  },
@@ -376,10 +345,8 @@ ObjectStorageResponse AzureObjStorageClient::delete_objects(const ObjectStorageP
 ObjectStorageResponse AzureObjStorageClient::delete_object(const ObjectStoragePathOptions& opts) {
  return do_azure_client_call(
  [&]() {
- auto resp = s3_put_rate_limit([&]() {
- SCOPED_BVAR_LATENCY(s3_bvar::s3_delete_object_latency);
- return _client->DeleteBlob(opts.key);
- });
+ SCOPED_BVAR_LATENCY(s3_bvar::s3_delete_object_latency);
+ auto resp = _client->DeleteBlob(opts.key);
  if (!resp.Value.Deleted) {
  throw Exception(Status::IOError<false>("Delete azure blob failed"));
  }
@@ -407,10 +374,8 @@ ObjectStorageResponse AzureObjStorageClient::delete_objects_recursively(
  ListBlobsPagedResponse resp;
  auto list_resp = do_azure_client_call(
  [&]() {
- resp = s3_get_rate_limit([&]() {
- SCOPED_BVAR_LATENCY(s3_bvar::s3_list_latency);
- return _client->ListBlobs(list_opts);
- });
+ SCOPED_BVAR_LATENCY(s3_bvar::s3_list_latency);
+ resp = _client->ListBlobs(list_opts);
  },
  opts, _tls_debug_context);
  if (list_resp.status.code != ErrorCode::OK) {
@@ -425,10 +390,8 @@ ObjectStorageResponse AzureObjStorageClient::delete_objects_recursively(
  list_opts.ContinuationToken = resp.NextPageToken;
  list_resp = do_azure_client_call(
  [&]() {
- resp = s3_get_rate_limit([&]() {
- SCOPED_BVAR_LATENCY(s3_bvar::s3_list_latency);
- return _client->ListBlobs(list_opts);
- });
+ SCOPED_BVAR_LATENCY(s3_bvar::s3_list_latency);
+ resp = _client->ListBlobs(list_opts);
  },
  opts, _tls_debug_context);
  if (list_resp.status.code != ErrorCode::OK) {