sollhui commented on code in PR #65420:
URL: https://github.com/apache/doris/pull/65420#discussion_r3592467477


##########
be/src/common/daemon.cpp:
##########
@@ -582,6 +584,12 @@ void Daemon::calculate_workload_group_metrics_thread() {
     }
 }
 
+void Daemon::refresh_cpu_info_thread() {
+    while 
(!_stop_background_threads_latch.wait_for(CPU_INFO_REFRESH_INTERVAL)) {
+        CpuInfo::refresh_current_num_cores();

Review Comment:
   Non-cloud BEs can still use the S3 QPS/bytes limiters, but the only 
production caller of `refresh_s3_rate_limiters()` is 
`CloudStorageEngine::_refresh_storage_vault_info_thread_callback()`. This 
generic daemon loop only refreshes the cached CPU count, and 
`S3ClientFactory::create()` no longer refreshes the holders. As a result, in 
classic mode mutable `s3_*` limiter config changes and CPU/cgroup changes never 
rebuild the active limiters and remain stale until restart. Please move the 
limiter refresh into a BE-wide periodic lifecycle (gated by 
`enable_s3_rate_limiter`) or add equivalent update callbacks, and cover the 
non-cloud lifecycle in a test.



##########
be/src/util/s3_util.cpp:
##########
@@ -156,67 +160,272 @@ 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";
+
+#ifdef BE_TEST
+std::atomic<int64_t> s3_rate_limiter_cpu_cores_for_test {0};
+#endif
+
+int64_t get_s3_rate_limiter_cpu_cores() {
+#ifdef BE_TEST
+    const int64_t cpu_cores_for_test =
+            s3_rate_limiter_cpu_cores_for_test.load(std::memory_order_relaxed);
+    if (cpu_cores_for_test > 0) {
+        return cpu_cores_for_test;
+    }
+#endif
+    return CpuInfo::get_current_num_cores();
+}
+
+int64_t get_s3_qps_per_core(S3RateLimitType type) {
+    switch (type) {
+    case S3RateLimitType::GET:
+        return config::s3_get_qps_per_core;
+    case S3RateLimitType::PUT:
+        return config::s3_put_qps_per_core;
+    default:
+        CHECK(false) << "unknown S3 rate limiter type: " << to_string(type);
+        return -1;
+    }
+}
+
+int64_t get_s3_bytes_per_second_per_core(S3RateLimitType type) {
+    switch (type) {
+    case S3RateLimitType::GET:
+        return config::s3_get_bytes_per_second_per_core;
+    case S3RateLimitType::PUT:
+        return config::s3_put_bytes_per_second_per_core;
+    default:
+        CHECK(false) << "unknown S3 rate limiter type: " << to_string(type);
+        return -1;
+    }
+}
+
+int64_t compute_s3_rate_limit_from_core(int64_t per_core, int64_t cores, 
int64_t max_rate) {
+    max_rate = max_rate > 0 ? max_rate : std::numeric_limits<int64_t>::max();
+    if (per_core > max_rate / cores) {
+        return max_rate;
+    }
+    return per_core * cores;
+}
+
+S3RateLimiterConfig 
get_effective_s3_rate_limiter_config_with_cpu_cores(S3RateLimitType type,
+                                                                        
int64_t cpu_cores) {
+    S3RateLimiterConfig result;
+    int64_t per_core;
+    int64_t max_qps;
+
+    switch (type) {
+    case S3RateLimitType::GET:
+        per_core = config::s3_get_qps_per_core;
+        max_qps = config::s3_get_qps_max;
+        if (per_core < 0) {
+            result.token_limit = config::s3_get_token_limit;
+            result.token_per_second = config::s3_get_token_per_second;
+            result.bucket_tokens = config::s3_get_bucket_tokens;
+        }
+        break;
+    case S3RateLimitType::PUT:
+        per_core = config::s3_put_qps_per_core;
+        max_qps = config::s3_put_qps_max;
+        if (per_core < 0) {
+            result.token_limit = config::s3_put_token_limit;
+            result.token_per_second = config::s3_put_token_per_second;
+            result.bucket_tokens = config::s3_put_bucket_tokens;
+        }
+        break;
+    default:
+        CHECK(false) << "unknown S3 rate limiter type: " << to_string(type);
+        return result;
+    }
+
+    if (per_core >= 0) {
+        result.token_per_second = compute_s3_rate_limit_from_core(per_core, 
cpu_cores, max_qps);
+        result.bucket_tokens = result.token_per_second;
+    }
+    return result;
+}
+
+S3RateLimiterConfig 
get_effective_s3_bytes_rate_limiter_config_with_cpu_cores(S3RateLimitType type,
+                                                                              
int64_t cpu_cores) {
+    int64_t per_core = 0;
+    int64_t max_bytes_per_second = 0;
+
+    switch (type) {
+    case S3RateLimitType::GET:
+        per_core = config::s3_get_bytes_per_second_per_core;
+        max_bytes_per_second = config::s3_get_bytes_per_second_max;
+        break;
+    case S3RateLimitType::PUT:
+        per_core = config::s3_put_bytes_per_second_per_core;
+        max_bytes_per_second = config::s3_put_bytes_per_second_max;
+        break;
+    default:
+        CHECK(false) << "unknown S3 rate limiter type: " << to_string(type);
+    }
+
+    S3RateLimiterConfig result;
+    if (per_core > 0) {
+        result.token_per_second =
+                compute_s3_rate_limit_from_core(per_core, cpu_cores, 
max_bytes_per_second);
+        result.bucket_tokens = result.token_per_second;
+    }
+    return result;
+}
 } // namespace
 
-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};
+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<bool> updating_get_limiter {false};
-static std::atomic<bool> updating_put_limiter {false};
+namespace {
+
+enum class S3RateLimiterKind {
+    QPS,
+    BYTES,
+};
+
+struct S3RateLimiterUpdateState {
+    S3RateLimiterConfig applied_config;
+    std::mutex update_lock;
+};
+
+std::array<std::array<S3RateLimiterUpdateState, 2>, 2> update_states;
+
+S3RateLimiterUpdateState& get_update_state(S3RateLimitType type, 
S3RateLimiterKind kind) {
+    DCHECK(type == S3RateLimitType::GET || type == S3RateLimitType::PUT);
+    return update_states[static_cast<size_t>(kind)][static_cast<size_t>(type)];
+}
+
+void update_s3_rate_limiter(S3ClientFactory& factory, S3RateLimitType type, 
S3RateLimiterKind kind,
+                            int64_t cpu_cores) {
+    auto& state = get_update_state(type, kind);
+    std::lock_guard update_guard(state.update_lock);
+
+    const int64_t per_core = kind == S3RateLimiterKind::BYTES
+                                     ? get_s3_bytes_per_second_per_core(type)
+                                     : get_s3_qps_per_core(type);
+    const auto current =
+            kind == S3RateLimiterKind::BYTES
+                    ? 
get_effective_s3_bytes_rate_limiter_config_with_cpu_cores(type, cpu_cores)
+                    : 
get_effective_s3_rate_limiter_config_with_cpu_cores(type, cpu_cores);
+    const bool config_changed = state.applied_config.token_per_second != 
current.token_per_second ||
+                                state.applied_config.bucket_tokens != 
current.bucket_tokens ||
+                                state.applied_config.token_limit != 
current.token_limit;
+    if (!config_changed) {
+        return;
+    }
+
+    auto* limiter = kind == S3RateLimiterKind::BYTES ? 
factory.bytes_rate_limiter(type)
+                                                     : 
factory.rate_limiter(type);
+    const int ret =
+            limiter->reset(current.token_per_second, current.bucket_tokens, 
current.token_limit);
+    if (ret != 0) {
+        LOG(WARNING) << "Failed to reset S3 " << to_string(type)
+                     << (kind == S3RateLimiterKind::BYTES ? " bytes" : " QPS")
+                     << " rate limiter, error code: " << ret;
+        return;
+    }
+
+    const int64_t logged_cpu_cores = per_core > 0 ? cpu_cores : 1;
+    if (kind == S3RateLimiterKind::BYTES) {
+        LOG(INFO) << "Reset S3 " << to_string(type)
+                  << " bytes rate limiter, bytes_per_second=" << 
current.token_per_second
+                  << ", burst_bytes=" << current.bucket_tokens
+                  << ", cpu_cores=" << logged_cpu_cores;
+    } else if (per_core >= 0) {
+        LOG(INFO) << "Reset S3 " << to_string(type)
+                  << " QPS rate limiter, qps=" << current.token_per_second
+                  << ", burst=" << current.bucket_tokens << ", limit=" << 
current.token_limit
+                  << ", cpu_cores=" << logged_cpu_cores;
+    }
+
+    state.applied_config = current;
+}
+
+} // namespace
 
 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();
 }
 
-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,
-                                    std::atomic<int64_t>& last_bucket,
-                                    std::atomic<int64_t>& last_limit,
-                                    std::atomic<bool>& updating_flag, const 
char* limiter_name) {
-    if (last_tps.load(std::memory_order_relaxed) != current_tps ||
-        last_bucket.load(std::memory_order_relaxed) != current_bucket ||
-        last_limit.load(std::memory_order_relaxed) != current_limit) {
-        bool expected = false;
-        if (!updating_flag.compare_exchange_strong(expected, true, 
std::memory_order_acq_rel)) {
-            return;
-        }
-        if (last_tps.load(std::memory_order_acquire) != current_tps ||
-            last_bucket.load(std::memory_order_acquire) != current_bucket ||
-            last_limit.load(std::memory_order_acquire) != current_limit) {
-            int ret =
-                    reset_s3_rate_limiter(LimiterType, current_tps, 
current_bucket, current_limit);
-
-            if (ret == 0) {
-                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);
-            } else {
-                LOG(WARNING) << "Failed to reset S3 " << limiter_name
-                             << " rate limiter, error code: " << ret;
-            }
-        }
+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();
+}
+
+bool acquire_s3_rate_limit(S3RateLimitType type, size_t bytes) {
+    if (!config::enable_s3_rate_limiter) {
+        return true;
+    }
 
-        updating_flag.store(false, std::memory_order_release);
+    auto& factory = S3ClientFactory::instance();
+    if (doris::apply_s3_rate_limit(type, factory.rate_limiter(type),
+                                   config::s3_rate_limiter_log_interval) < 0) {
+        return false;
     }
+    if (bytes == 0) {
+        return true;
+    }
+    return factory.bytes_rate_limiter(type)->add(bytes) >= 0;

Review Comment:
   The bytes configs default to `-1`, which produces a disabled holder 
(`max_speed=0`, `limit=0`), but every data request with `bytes > 0` still calls 
`add(bytes)`. `add_with_config()` takes the holder's shared mutex and 
`TokenBucketRateLimiter::add()` then takes its spinlock and increments 
`_count`, so enabling only the QPS limiter adds global synchronization to every 
S3/Azure data request for a no-op bytes limiter. Please add a fast path based 
on the active bytes-limiter config (not merely the mutable config value, so 
periodic-refresh semantics are preserved).



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to