be/src/util/s3_rate_limiter_manager.cpp

Source
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#include "util/s3_rate_limiter_manager.h"

#include <algorithm>
#include <limits>
#include <thread>

#include "common/config.h"
#include "common/logging.h"
#include "util/cgroup_util.h"

namespace doris {

bvar::Adder<int64_t> s3_get_bytes_rate_limit_sleep_ns("s3_get_bytes_rate_limit_sleep_ns");
bvar::Adder<int64_t> s3_get_bytes_rate_limit_sleep_count("s3_get_bytes_rate_limit_sleep_count");
bvar::Adder<int64_t> s3_get_bytes_rate_limit_rejected_count(
        "s3_get_bytes_rate_limit_rejected_count");
bvar::Adder<int64_t> s3_put_bytes_rate_limit_sleep_ns("s3_put_bytes_rate_limit_sleep_ns");
bvar::Adder<int64_t> s3_put_bytes_rate_limit_sleep_count("s3_put_bytes_rate_limit_sleep_count");
bvar::Adder<int64_t> s3_put_bytes_rate_limit_rejected_count(
        "s3_put_bytes_rate_limit_rejected_count");

namespace {

std::function<void(int64_t)> bytes_rate_limiter_metric_func(S3RateLimitType type) {
    switch (type) {
    case S3RateLimitType::GET:
        return metric_func_factory(s3_get_bytes_rate_limit_sleep_ns,
                                   s3_get_bytes_rate_limit_sleep_count,
                                   &s3_get_bytes_rate_limit_rejected_count);
    case S3RateLimitType::PUT:
        return metric_func_factory(s3_put_bytes_rate_limit_sleep_ns,
                                   s3_put_bytes_rate_limit_sleep_count,
                                   &s3_put_bytes_rate_limit_rejected_count);
    default:
        return [](int64_t) {};
    }
}

// min(per_core * cores, cap) with overflow protection; cap == 0 means no cap.
int64_t cap_multiply(int64_t per_core, int64_t cores, int64_t cap) {
    cap = cap > 0 ? cap : std::numeric_limits<int64_t>::max();
    if (per_core > cap / cores) {
        return cap;
    }
    return per_core * cores;
}

size_t index_of(S3RateLimitType type) {
    CHECK(type == S3RateLimitType::GET || type == S3RateLimitType::PUT) << to_string(type);
    return static_cast<size_t>(type);
}

} // namespace

S3EffectiveRateLimit resolve_s3_rate_limit(S3RateLimitType type, int64_t cores) {
    const bool is_get = type == S3RateLimitType::GET;
    const int64_t qps_per_core = is_get ? config::s3_get_qps_per_core : config::s3_put_qps_per_core;
    const int64_t qps_max = is_get ? config::s3_get_qps_max : config::s3_put_qps_max;
    const int64_t bytes_per_core = is_get ? config::s3_get_bytes_per_second_per_core
                                          : config::s3_put_bytes_per_second_per_core;
    const int64_t bytes_max =
            is_get ? config::s3_get_bytes_per_second_max : config::s3_put_bytes_per_second_max;
    cores = std::max<int64_t>(1, cores);

    S3EffectiveRateLimit limit;
    if (qps_per_core < 0) {
        // Unset: the legacy absolute configs stay in charge, bit-for-bit compatible.
        limit.qps = is_get ? config::s3_get_token_per_second : config::s3_put_token_per_second;
        limit.burst = is_get ? config::s3_get_bucket_tokens : config::s3_put_bucket_tokens;
        limit.count_limit = is_get ? config::s3_get_token_limit : config::s3_put_token_limit;
    } else if (qps_per_core > 0) {
        limit.qps = cap_multiply(qps_per_core, cores, qps_max);
        limit.burst = limit.qps; // burst = 1 second worth of quota
    }                            // qps_per_core == 0: QPS limiting disabled, all fields stay 0.

    if (bytes_per_core > 0) {
        limit.bytes_per_second = cap_multiply(bytes_per_core, cores, bytes_max);
    }
    return limit;
}

int reset_s3_rate_limiter(S3RateLimitType type, size_t max_speed, size_t max_burst, size_t limit) {
    if (type == S3RateLimitType::UNKNOWN) {
        return -1;
    }
    return S3RateLimiterManager::instance().qps_limiter(type)->reset(max_speed, max_burst, limit);
}

int64_t s3_rate_limiter_cpu_cores() {
    if (int64_t overridden = config::s3_rate_limiter_cpu_cores; overridden > 0) {
        return overridden;
    }
    int physical = static_cast<int>(std::thread::hardware_concurrency());
    // Re-read the cgroup quota on every call: serverless BEs can be resized in place,
    // and the daemon refresh thread picks the change up through here.
    int limited = CGroupUtil::get_cgroup_limited_cpu_number(physical);
    return std::max(1, limited);
}

S3RateLimiterManager::S3RateLimiterManager() {
    const int64_t cores = s3_rate_limiter_cpu_cores();
    for (auto type : {S3RateLimitType::GET, S3RateLimitType::PUT}) {
        auto limit = resolve_s3_rate_limit(type, cores);
        _qps_limiters[index_of(type)] = std::make_unique<S3RateLimiterHolder>(
                limit.qps, limit.burst, limit.count_limit, s3_rate_limiter_metric_func(type));
        _bytes_limiters[index_of(type)] = std::make_unique<S3RateLimiterHolder>(
                limit.bytes_per_second, limit.bytes_per_second, 0,
                bytes_rate_limiter_metric_func(type));
    }
}

S3RateLimiterManager& S3RateLimiterManager::instance() {
    static S3RateLimiterManager ret;
    return ret;
}

S3RateLimiterHolder* S3RateLimiterManager::qps_limiter(S3RateLimitType type) {
    return _qps_limiters[index_of(type)].get();
}

S3RateLimiterHolder* S3RateLimiterManager::bytes_limiter(S3RateLimitType type) {
    return _bytes_limiters[index_of(type)].get();
}

void S3RateLimiterManager::refresh() {
    std::lock_guard guard(_refresh_lock);
    const int64_t cores = s3_rate_limiter_cpu_cores();
    for (auto type : {S3RateLimitType::GET, S3RateLimitType::PUT}) {
        const auto limit = resolve_s3_rate_limit(type, cores);

        auto* qps = qps_limiter(type);
        if (qps->get_max_speed() != static_cast<size_t>(limit.qps) ||
            qps->get_max_burst() != static_cast<size_t>(limit.burst) ||
            qps->get_limit() != static_cast<size_t>(limit.count_limit)) {
            qps->reset(limit.qps, limit.burst, limit.count_limit);
            LOG(INFO) << "reset S3 " << to_string(type) << " QPS rate limiter, qps=" << limit.qps
                      << ", burst=" << limit.burst << ", count_limit=" << limit.count_limit
                      << ", cores=" << cores;
        }

        auto* bytes = bytes_limiter(type);
        if (bytes->get_max_speed() != static_cast<size_t>(limit.bytes_per_second)) {
            bytes->reset(limit.bytes_per_second, limit.bytes_per_second, 0);
            LOG(INFO) << "reset S3 " << to_string(type)
                      << " bytes rate limiter, bytes_per_second=" << limit.bytes_per_second
                      << ", cores=" << cores;
        }
    }
}

S3RateLimitGuard::S3RateLimitGuard(S3RateLimitType type, size_t estimated_bytes) {
    if (!config::enable_s3_rate_limiter) {
        return;
    }
    auto& mgr = S3RateLimiterManager::instance();

    auto* qps = mgr.qps_limiter(type);
    if (qps->is_enabled() &&
        apply_s3_rate_limit(type, qps, config::s3_rate_limiter_log_interval) < 0) {
        _ok = false;
        _reject_reason = S3RateLimitRejectReason::QPS;
        return;
    }

    if (estimated_bytes == 0) {
        return;
    }
    auto* bytes = mgr.bytes_limiter(type);
    if (!bytes->is_enabled()) {
        return;
    }
    // Clamp the reservation to 1 second worth of bandwidth so a single oversized IO
    // (e.g. a whole-file read_at) cannot create unbounded upfront debt. The clamped
    // remainder is intentionally not accounted; effective quotas below the single-IO
    // upper bound are excluded by the config contract (see config.cpp).
    _reserved = std::min(estimated_bytes, bytes->get_max_speed());
    if (_reserved > 0) {
        // Pin the admitted bucket generation for settle().
        _charged_bucket = bytes->charge(_reserved);
        if (_charged_bucket == nullptr) {
            _ok = false;
            _reject_reason = S3RateLimitRejectReason::BYTES;
        }
    }
}

void S3RateLimitGuard::settle(size_t actual_bytes) {
    if (_settled) {
        return;
    }
    _settled = true;
    if (_charged_bucket != nullptr && _reserved > actual_bytes) {
        _charged_bucket->refund(_reserved - actual_bytes);
    }
}

} // namespace doris

Coverage Report

Created: 2026-07-20 21:29

Line	Count	Source
1		// Licensed to the Apache Software Foundation (ASF) under one
2		// or more contributor license agreements. See the NOTICE file
3		// distributed with this work for additional information
4		// regarding copyright ownership. The ASF licenses this file
5		// to you under the Apache License, Version 2.0 (the
6		// "License"); you may not use this file except in compliance
7		// with the License. You may obtain a copy of the License at
8		//
9		// http://www.apache.org/licenses/LICENSE-2.0
10		//
11		// Unless required by applicable law or agreed to in writing,
12		// software distributed under the License is distributed on an
13		// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14		// KIND, either express or implied. See the License for the
15		// specific language governing permissions and limitations
16		// under the License.
17
18		#include "util/s3_rate_limiter_manager.h"
19
20		#include <algorithm>
21		#include <limits>
22		#include <thread>
23
24		#include "common/config.h"
25		#include "common/logging.h"
26		#include "util/cgroup_util.h"
27
28		namespace doris {
29
30		bvar::Adder<int64_t> s3_get_bytes_rate_limit_sleep_ns("s3_get_bytes_rate_limit_sleep_ns");
31		bvar::Adder<int64_t> s3_get_bytes_rate_limit_sleep_count("s3_get_bytes_rate_limit_sleep_count");
32		bvar::Adder<int64_t> s3_get_bytes_rate_limit_rejected_count(
33		"s3_get_bytes_rate_limit_rejected_count");
34		bvar::Adder<int64_t> s3_put_bytes_rate_limit_sleep_ns("s3_put_bytes_rate_limit_sleep_ns");
35		bvar::Adder<int64_t> s3_put_bytes_rate_limit_sleep_count("s3_put_bytes_rate_limit_sleep_count");
36		bvar::Adder<int64_t> s3_put_bytes_rate_limit_rejected_count(
37		"s3_put_bytes_rate_limit_rejected_count");
38
39		namespace {
40
41	2	std::function<void(int64_t)> bytes_rate_limiter_metric_func(S3RateLimitType type) {
42	2	switch (type) {
43	1	case S3RateLimitType::GET:
44	1	return metric_func_factory(s3_get_bytes_rate_limit_sleep_ns,
45	1	s3_get_bytes_rate_limit_sleep_count,
46	1	&s3_get_bytes_rate_limit_rejected_count);
47	1	case S3RateLimitType::PUT:
48	1	return metric_func_factory(s3_put_bytes_rate_limit_sleep_ns,
49	1	s3_put_bytes_rate_limit_sleep_count,
50	1	&s3_put_bytes_rate_limit_rejected_count);
51	0	default:
52	0	return [](int64_t) {};
53	2	}
54	2	}
55
56		// min(per_core * cores, cap) with overflow protection; cap == 0 means no cap.
57	9	int64_t cap_multiply(int64_t per_core, int64_t cores, int64_t cap) {
58	9	cap = cap > 0 ? cap : std::numeric_limits<int64_t>::max();
59	9	if (per_core > cap / cores) {
60	4	return cap;
61	4	}
62	5	return per_core * cores;
63	9	}
64
65	213	size_t index_of(S3RateLimitType type) {
66	213	CHECK(type == S3RateLimitType::GET \|\| type == S3RateLimitType::PUT) << to_string(type);
67	213	return static_cast<size_t>(type);
68	213	}
69
70		} // namespace
71
72	83	S3EffectiveRateLimit resolve_s3_rate_limit(S3RateLimitType type, int64_t cores) {
73	83	const bool is_get = type == S3RateLimitType::GET;
74	83	const int64_t qps_per_core = is_get ? config::s3_get_qps_per_core : config::s3_put_qps_per_core;
75	83	const int64_t qps_max = is_get ? config::s3_get_qps_max : config::s3_put_qps_max;
76	83	const int64_t bytes_per_core = is_get ? config::s3_get_bytes_per_second_per_core
77	83	: config::s3_put_bytes_per_second_per_core;
78	83	const int64_t bytes_max =
79	83	is_get ? config::s3_get_bytes_per_second_max : config::s3_put_bytes_per_second_max;
80	83	cores = std::max<int64_t>(1, cores);
81
82	83	S3EffectiveRateLimit limit;
83	83	if (qps_per_core < 0) {
84		// Unset: the legacy absolute configs stay in charge, bit-for-bit compatible.
85	76	limit.qps = is_get ? config::s3_get_token_per_second : config::s3_put_token_per_second;
86	76	limit.burst = is_get ? config::s3_get_bucket_tokens : config::s3_put_bucket_tokens;
87	76	limit.count_limit = is_get ? config::s3_get_token_limit : config::s3_put_token_limit;
88	76	} else if (qps_per_core > 0) {
89	6	limit.qps = cap_multiply(qps_per_core, cores, qps_max);
90	6	limit.burst = limit.qps; // burst = 1 second worth of quota
91	6	} // qps_per_core == 0: QPS limiting disabled, all fields stay 0.
92
93	83	if (bytes_per_core > 0) {
94	3	limit.bytes_per_second = cap_multiply(bytes_per_core, cores, bytes_max);
95	3	}
96	83	return limit;
97	83	}
98
99	0	int reset_s3_rate_limiter(S3RateLimitType type, size_t max_speed, size_t max_burst, size_t limit) {
100	0	if (type == S3RateLimitType::UNKNOWN) {
101	0	return -1;
102	0	}
103	0	return S3RateLimiterManager::instance().qps_limiter(type)->reset(max_speed, max_burst, limit);
104	0	}
105
106	39	int64_t s3_rate_limiter_cpu_cores() {
107	39	if (int64_t overridden = config::s3_rate_limiter_cpu_cores; overridden > 0) {
108	6	return overridden;
109	6	}
110	33	int physical = static_cast<int>(std::thread::hardware_concurrency());
111		// Re-read the cgroup quota on every call: serverless BEs can be resized in place,
112		// and the daemon refresh thread picks the change up through here.
113	33	int limited = CGroupUtil::get_cgroup_limited_cpu_number(physical);
114	33	return std::max(1, limited);
115	39	}
116
117	1	S3RateLimiterManager::S3RateLimiterManager() {
118	1	const int64_t cores = s3_rate_limiter_cpu_cores();
119	2	for (auto type : {S3RateLimitType::GET, S3RateLimitType::PUT}) {
120	2	auto limit = resolve_s3_rate_limit(type, cores);
121	2	_qps_limiters[index_of(type)] = std::make_unique<S3RateLimiterHolder>(
122	2	limit.qps, limit.burst, limit.count_limit, s3_rate_limiter_metric_func(type));
123	2	_bytes_limiters[index_of(type)] = std::make_unique<S3RateLimiterHolder>(
124	2	limit.bytes_per_second, limit.bytes_per_second, 0,
125	2	bytes_rate_limiter_metric_func(type));
126	2	}
127	1	}
128
129	64	S3RateLimiterManager& S3RateLimiterManager::instance() {
130	64	static S3RateLimiterManager ret;
131	64	return ret;
132	64	}
133
134	114	S3RateLimiterHolder* S3RateLimiterManager::qps_limiter(S3RateLimitType type) {
135	114	return _qps_limiters[index_of(type)].get();
136	114	}
137
138	95	S3RateLimiterHolder* S3RateLimiterManager::bytes_limiter(S3RateLimitType type) {
139	95	return _bytes_limiters[index_of(type)].get();
140	95	}
141
142	36	void S3RateLimiterManager::refresh() {
143	36	std::lock_guard guard(_refresh_lock);
144	36	const int64_t cores = s3_rate_limiter_cpu_cores();
145	72	for (auto type : {S3RateLimitType::GET, S3RateLimitType::PUT}) {
146	72	const auto limit = resolve_s3_rate_limit(type, cores);
147
148	72	auto* qps = qps_limiter(type);
149	72	if (qps->get_max_speed() != static_cast<size_t>(limit.qps) \|\|
150	72	qps->get_max_burst() != static_cast<size_t>(limit.burst) \|\|
151	72	qps->get_limit() != static_cast<size_t>(limit.count_limit)) {
152	27	qps->reset(limit.qps, limit.burst, limit.count_limit);
153	27	LOG(INFO) << "reset S3 " << to_string(type) << " QPS rate limiter, qps=" << limit.qps
154	27	<< ", burst=" << limit.burst << ", count_limit=" << limit.count_limit
155	27	<< ", cores=" << cores;
156	27	}
157
158	72	auto* bytes = bytes_limiter(type);
159	72	if (bytes->get_max_speed() != static_cast<size_t>(limit.bytes_per_second)) {
160	10	bytes->reset(limit.bytes_per_second, limit.bytes_per_second, 0);
161	10	LOG(INFO) << "reset S3 " << to_string(type)
162	10	<< " bytes rate limiter, bytes_per_second=" << limit.bytes_per_second
163	10	<< ", cores=" << cores;
164	10	}
165	72	}
166	36	}
167
168	237k	S3RateLimitGuard::S3RateLimitGuard(S3RateLimitType type, size_t estimated_bytes) {
169	237k	if (!config::enable_s3_rate_limiter) {
170	237k	return;
171	237k	}
172	34	auto& mgr = S3RateLimiterManager::instance();
173
174	34	auto* qps = mgr.qps_limiter(type);
175	34	if (qps->is_enabled() &&
176	34	apply_s3_rate_limit(type, qps, config::s3_rate_limiter_log_interval) < 0) {
177	5	_ok = false;
178	5	_reject_reason = S3RateLimitRejectReason::QPS;
179	5	return;
180	5	}
181
182	29	if (estimated_bytes == 0) {
183	7	return;
184	7	}
185	22	auto* bytes = mgr.bytes_limiter(type);
186	22	if (!bytes->is_enabled()) {
187	1	return;
188	1	}
189		// Clamp the reservation to 1 second worth of bandwidth so a single oversized IO
190		// (e.g. a whole-file read_at) cannot create unbounded upfront debt. The clamped
191		// remainder is intentionally not accounted; effective quotas below the single-IO
192		// upper bound are excluded by the config contract (see config.cpp).
193	21	_reserved = std::min(estimated_bytes, bytes->get_max_speed());
194	21	if (_reserved > 0) {
195		// Pin the admitted bucket generation for settle().
196	8	_charged_bucket = bytes->charge(_reserved);
197	8	if (_charged_bucket == nullptr) {
198	2	_ok = false;
199	2	_reject_reason = S3RateLimitRejectReason::BYTES;
200	2	}
201	8	}
202	21	}
203
204	188k	void S3RateLimitGuard::settle(size_t actual_bytes) {
205	188k	if (_settled) {
206	0	return;
207	0	}
208	188k	_settled = true;
209	188k	if (_charged_bucket != nullptr && _reserved > actual_bytes) {
210	3	_charged_bucket->refund(_reserved - actual_bytes);
211	3	}
212	188k	}
213
214		} // namespace doris