mirror of
https://github.com/apache/doris.git
synced 2026-11-04 17:52:47 +00:00
[feature](be) Add file range read scheduler
### What problem does this PR solve? Issue Number: #67611 (tracking) Related PR: #67292, #67579, #67612 Problem Summary: Provide query/BE resident-byte admission, one exact range read per existing thread-pool task, cancellation, and owned buffer handles. Keep terminal-state checks on FileRangeRead and rejected-result construction on FileRangeReadSubmitResult. Preserve accepted-task accounting and atomic shared S3 reader statistics for concurrent requests. Admit each planned range that fits the remaining query and BE budgets in input order. Preserve the original result slots for skipped ranges and allocate only admitted buffers, so wide scans retain partial prefetch instead of rejecting an entire batch. Budget checks remain one linear pass under both locks; allocation failure and cancellation release all reservations. ### Release note Internal asynchronous file-range reads with query/BE byte budgets and IO accounting. ### Check List (For Author) - Test: Unit Test — integrated ASAN/PCH validation covers 141 tests from 8 suites through ./run-be-ut.sh --run -j100. Includes partial budget admission, original page offsets, cancellation, allocation rollback, executor rejection, consumed-range callbacks and row-store switch fallback. Individual intermediate commits are not built separately; target-branch validation is recorded on the PR. Clang-format 16, header hygiene and diff checks pass. No cluster or clang-tidy run. - Behavior changed: Yes — introduces the opt-in component; this synchronization preserves existing defaults and IO/writeback semantics. - Does this need documentation: Yes — component scope and current behavior are tracked in #67611 and the related PR descriptions.
This commit is contained in:
1 parent
90cf5b6ce6
commit
85894d0c6e
7 files changed
+2165
-6
No files matched your search
@@ -0,0 +1,663 @@
|
||||
// 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 "io/fs/file_range_read_scheduler.h"
|
||||
|
||||
#include <exception>
|
||||
#include <limits>
|
||||
#include <utility>
|
||||
|
||||
#include "common/exception.h"
|
||||
#include "common/logging.h"
|
||||
#include "core/allocator.h"
|
||||
#include "cpp/sync_point.h"
|
||||
#include "io/fs/read_ahead_metrics.h"
|
||||
#include "runtime/memory/mem_tracker_limiter.h"
|
||||
#include "runtime/thread_context.h"
|
||||
#include "util/defer_op.h"
|
||||
#include "util/threadpool.h"
|
||||
#include "util/time.h"
|
||||
|
||||
namespace doris::io {
|
||||
|
||||
using FileRangeReadAllocator = Allocator<false, false, false, DefaultMemoryAllocator, true>;
|
||||
|
||||
namespace {
|
||||
|
||||
bool add_overflows(size_t left, size_t right) {
|
||||
return right > std::numeric_limits<size_t>::max() - left;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
FileRangeReadSubmitResult FileRangeReadSubmitResult::rejected(FileRangeReadRejectReason reason,
|
||||
Status status) {
|
||||
return FileRangeReadSubmitResult {
|
||||
.reads = {},
|
||||
.reject_reason = reason,
|
||||
.status = std::move(status),
|
||||
};
|
||||
}
|
||||
|
||||
FileRangeReadBudget::FileRangeReadBudget(size_t max_bytes) : _max_bytes(max_bytes) {
|
||||
DORIS_CHECK(_max_bytes > 0);
|
||||
}
|
||||
|
||||
size_t FileRangeReadBudget::resident_bytes() const {
|
||||
std::lock_guard lock(_mutex);
|
||||
return _resident_bytes;
|
||||
}
|
||||
|
||||
bool FileRangeReadBudget::_can_reserve_unlocked(size_t bytes) const {
|
||||
return bytes <= _max_bytes - _resident_bytes;
|
||||
}
|
||||
|
||||
void FileRangeReadBudget::_reserve_unlocked(size_t bytes) {
|
||||
DORIS_CHECK(_can_reserve_unlocked(bytes));
|
||||
_resident_bytes += bytes;
|
||||
}
|
||||
|
||||
void FileRangeReadBudget::_release(size_t bytes) {
|
||||
std::lock_guard lock(_mutex);
|
||||
DORIS_CHECK(bytes <= _resident_bytes);
|
||||
_resident_bytes -= bytes;
|
||||
}
|
||||
|
||||
FileRangeReadContext::FileRangeReadContext(size_t max_bytes) : _budget(max_bytes) {}
|
||||
|
||||
void FileRangeReadContext::cancel() {
|
||||
if (_cancelled.exchange(true, std::memory_order_acq_rel)) {
|
||||
return;
|
||||
}
|
||||
|
||||
std::lock_guard lock(_reads_mutex);
|
||||
for (const auto& weak_read : _reads) {
|
||||
if (auto read = weak_read.lock(); read != nullptr) {
|
||||
read->request_cancel();
|
||||
}
|
||||
}
|
||||
_reads.clear();
|
||||
}
|
||||
|
||||
void FileRangeReadContext::_register_reads(
|
||||
const std::vector<std::shared_ptr<FileRangeRead>>& reads) {
|
||||
bool cancel_reads = false;
|
||||
{
|
||||
std::lock_guard lock(_reads_mutex);
|
||||
cancel_reads = _cancelled.load(std::memory_order_acquire);
|
||||
if (!cancel_reads) {
|
||||
std::erase_if(_reads, [](const auto& read) { return read.expired(); });
|
||||
_reads.reserve(_reads.size() + reads.size());
|
||||
for (const auto& read : reads) {
|
||||
// Partial admission leaves null slots for ranges without a buffer or task.
|
||||
if (read != nullptr) {
|
||||
_reads.emplace_back(read);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if (cancel_reads) {
|
||||
for (const auto& read : reads) {
|
||||
if (read != nullptr) {
|
||||
read->request_cancel();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
FileRangeReadReservation::FileRangeReadReservation(
|
||||
std::shared_ptr<FileRangeReadContext> context,
|
||||
std::shared_ptr<FileRangeReadBudget> global_budget, size_t bytes) noexcept
|
||||
: _context(std::move(context)), _global_budget(std::move(global_budget)), _bytes(bytes) {
|
||||
DORIS_CHECK(_context != nullptr);
|
||||
DORIS_CHECK(_global_budget != nullptr);
|
||||
DORIS_CHECK(_bytes > 0);
|
||||
}
|
||||
|
||||
FileRangeReadReservation::FileRangeReadReservation(FileRangeReadReservation&& other) noexcept
|
||||
: _context(std::move(other._context)),
|
||||
_global_budget(std::move(other._global_budget)),
|
||||
_bytes(std::exchange(other._bytes, 0)) {}
|
||||
|
||||
FileRangeReadReservation& FileRangeReadReservation::operator=(
|
||||
FileRangeReadReservation&& other) noexcept {
|
||||
if (this != &other) {
|
||||
_reset();
|
||||
_context = std::move(other._context);
|
||||
_global_budget = std::move(other._global_budget);
|
||||
_bytes = std::exchange(other._bytes, 0);
|
||||
}
|
||||
return *this;
|
||||
}
|
||||
|
||||
FileRangeReadReservation::~FileRangeReadReservation() {
|
||||
_reset();
|
||||
}
|
||||
|
||||
void FileRangeReadReservation::_reset() {
|
||||
if (_context == nullptr) {
|
||||
return;
|
||||
}
|
||||
_context->_budget._release(_bytes);
|
||||
_global_budget->_release(_bytes);
|
||||
read_ahead_bvars().resident_bytes << -static_cast<int64_t>(_bytes);
|
||||
_context.reset();
|
||||
_global_budget.reset();
|
||||
_bytes = 0;
|
||||
}
|
||||
|
||||
FileRangeRead::FileRangeRead(FileRange range, std::shared_ptr<MemTrackerLimiter> tracker,
|
||||
FileRangeReadReservation reservation,
|
||||
std::shared_ptr<ReadAheadStatistics> statistics)
|
||||
: _range(range),
|
||||
_tracker(std::move(tracker)),
|
||||
_reservation(std::move(reservation)),
|
||||
_statistics(std::move(statistics)) {
|
||||
FileRangeReadAllocator allocator;
|
||||
_data = reinterpret_cast<char*>(allocator.alloc(_range.size));
|
||||
}
|
||||
|
||||
FileRangeRead::~FileRangeRead() {
|
||||
SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(_tracker);
|
||||
FileRangeReadAllocator allocator;
|
||||
allocator.free(_data, _range.size);
|
||||
}
|
||||
|
||||
Status FileRangeRead::create(FileRange range, std::shared_ptr<MemTrackerLimiter> tracker,
|
||||
FileRangeReadReservation reservation,
|
||||
std::shared_ptr<FileRangeRead>* output,
|
||||
std::shared_ptr<ReadAheadStatistics> statistics) {
|
||||
DORIS_CHECK(range.size > 0);
|
||||
DORIS_CHECK(tracker != nullptr);
|
||||
DORIS_CHECK(output != nullptr);
|
||||
|
||||
Status injected_status;
|
||||
TEST_SYNC_POINT_CALLBACK("FileRangeRead::create:inject_failure", &injected_status);
|
||||
if (!injected_status.ok()) {
|
||||
return injected_status;
|
||||
}
|
||||
|
||||
SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(tracker);
|
||||
try {
|
||||
output->reset(new FileRangeRead(range, std::move(tracker), std::move(reservation),
|
||||
std::move(statistics)));
|
||||
} catch (const Exception& exception) {
|
||||
return exception.to_status();
|
||||
} catch (const std::exception& exception) {
|
||||
return Status::MemoryAllocFailed("allocate file range read buffer failed: {}",
|
||||
exception.what());
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
bool FileRangeRead::_is_terminal(State state) {
|
||||
return state == State::READY || state == State::FAILED || state == State::CANCELLED;
|
||||
}
|
||||
|
||||
Status FileRangeRead::wait(bool* ready_on_entry) {
|
||||
std::unique_lock lock(_mutex);
|
||||
if (ready_on_entry != nullptr) {
|
||||
*ready_on_entry = _state == State::READY;
|
||||
}
|
||||
_cv.wait(lock, [this]() { return _is_terminal(_state); });
|
||||
return _status;
|
||||
}
|
||||
|
||||
void FileRangeRead::request_cancel() {
|
||||
bool publish_cancelled = false;
|
||||
{
|
||||
std::lock_guard lock(_mutex);
|
||||
if (_is_terminal(_state)) {
|
||||
return;
|
||||
}
|
||||
_cancel_requested = true;
|
||||
publish_cancelled = _state == State::QUEUED;
|
||||
if (publish_cancelled) {
|
||||
_state = State::CANCELLED;
|
||||
_status = Status::Cancelled("asynchronous file range read cancelled");
|
||||
}
|
||||
}
|
||||
if (publish_cancelled) {
|
||||
_cv.notify_all();
|
||||
}
|
||||
}
|
||||
|
||||
FileRangeRead::State FileRangeRead::state() const {
|
||||
std::lock_guard lock(_mutex);
|
||||
return _state;
|
||||
}
|
||||
|
||||
Slice FileRangeRead::data() const {
|
||||
return slice(0, _range.size);
|
||||
}
|
||||
|
||||
Slice FileRangeRead::slice(size_t buffer_offset, size_t size) const {
|
||||
std::lock_guard lock(_mutex);
|
||||
DORIS_CHECK(_state == State::READY);
|
||||
DORIS_CHECK(buffer_offset <= _range.size);
|
||||
DORIS_CHECK(size <= _range.size - buffer_offset);
|
||||
return {_data + buffer_offset, size};
|
||||
}
|
||||
|
||||
FileRangeReadStats FileRangeRead::stats() const {
|
||||
std::lock_guard lock(_mutex);
|
||||
return _stats;
|
||||
}
|
||||
|
||||
bool FileRangeRead::_mark_running() {
|
||||
std::lock_guard lock(_mutex);
|
||||
if (_state == State::CANCELLED) {
|
||||
return false;
|
||||
}
|
||||
DORIS_CHECK(_state == State::QUEUED);
|
||||
_state = State::RUNNING;
|
||||
return true;
|
||||
}
|
||||
|
||||
bool FileRangeRead::_is_cancel_requested() const {
|
||||
std::lock_guard lock(_mutex);
|
||||
return _cancel_requested;
|
||||
}
|
||||
|
||||
void FileRangeRead::_publish_ready(FileRangeReadStats stats) {
|
||||
_publish_from_running(State::READY, Status::OK(), std::move(stats));
|
||||
}
|
||||
|
||||
void FileRangeRead::_publish_failed(Status status, FileRangeReadStats stats) {
|
||||
DORIS_CHECK(!status.ok());
|
||||
_publish_from_running(State::FAILED, std::move(status), std::move(stats));
|
||||
}
|
||||
|
||||
void FileRangeRead::_publish_submit_failure(Status status) {
|
||||
DORIS_CHECK(!status.ok());
|
||||
{
|
||||
std::lock_guard lock(_mutex);
|
||||
if (_state == State::CANCELLED) {
|
||||
return;
|
||||
}
|
||||
DORIS_CHECK(_state == State::QUEUED);
|
||||
_state = State::FAILED;
|
||||
_status = std::move(status);
|
||||
}
|
||||
_cv.notify_all();
|
||||
}
|
||||
|
||||
void FileRangeRead::_publish_cancelled(FileRangeReadStats stats) {
|
||||
_publish_from_running(State::CANCELLED,
|
||||
Status::Cancelled("asynchronous file range read cancelled"),
|
||||
std::move(stats));
|
||||
}
|
||||
|
||||
void FileRangeRead::_publish_from_running(State state, Status status, FileRangeReadStats stats) {
|
||||
{
|
||||
std::lock_guard lock(_mutex);
|
||||
DORIS_CHECK(_state == State::RUNNING);
|
||||
DORIS_CHECK(_is_terminal(state));
|
||||
if (_cancel_requested) {
|
||||
state = State::CANCELLED;
|
||||
status = Status::Cancelled("asynchronous file range read cancelled");
|
||||
}
|
||||
_state = state;
|
||||
_status = std::move(status);
|
||||
_stats = std::move(stats);
|
||||
// Publish accounting before waking a consumer that may immediately close its scanner.
|
||||
if (_statistics != nullptr) {
|
||||
auto& outcome = state == State::READY ? _statistics->completed_ranges
|
||||
: state == State::FAILED ? _statistics->failed_ranges
|
||||
: _statistics->cancelled_ranges;
|
||||
COUNTER_UPDATE(&outcome, 1);
|
||||
COUNTER_UPDATE(&_statistics->queue_wait_time, _stats.queue_wait_ns);
|
||||
COUNTER_UPDATE(&_statistics->read_time, _stats.read_ns);
|
||||
COUNTER_UPDATE(&_statistics->source_bytes, static_cast<int64_t>(_stats.bytes_read));
|
||||
COUNTER_UPDATE(&_statistics->successful_source_bytes,
|
||||
static_cast<int64_t>(_stats.successful_source_bytes));
|
||||
COUNTER_UPDATE(&_statistics->remote_bytes, _stats.file_cache.bytes_read_from_remote);
|
||||
COUNTER_UPDATE(&_statistics->local_bytes, _stats.file_cache.bytes_read_from_local);
|
||||
COUNTER_UPDATE(&_statistics->remote_requests, _stats.file_cache.num_remote_io_total);
|
||||
}
|
||||
auto& metrics = read_ahead_bvars();
|
||||
metrics.source_bytes << static_cast<int64_t>(_stats.bytes_read);
|
||||
metrics.successful_source_bytes << static_cast<int64_t>(_stats.successful_source_bytes);
|
||||
metrics.remote_bytes << _stats.file_cache.bytes_read_from_remote;
|
||||
metrics.remote_requests << _stats.file_cache.num_remote_io_total;
|
||||
if (state == State::FAILED) {
|
||||
read_ahead_bvars().failed_ranges << 1;
|
||||
}
|
||||
}
|
||||
_cv.notify_all();
|
||||
}
|
||||
|
||||
FileRangeReadIOContext FileRangeReadIOContext::from_caller(const IOContext& source) {
|
||||
FileRangeReadIOContext result;
|
||||
result.io_context = source;
|
||||
if (source.query_id != nullptr) {
|
||||
result.query_id = std::make_shared<TUniqueId>(*source.query_id);
|
||||
}
|
||||
result.io_context.file_cache_stats = nullptr;
|
||||
result.io_context.file_reader_stats = nullptr;
|
||||
result.io_context.remote_scan_cache_write_limiter = nullptr;
|
||||
result.io_context.condition_cache_filtered_rows = 0;
|
||||
result.io_context.predicate_filtered_rows = 0;
|
||||
result.io_context.cache_write_mode_override = CacheWriteMode::NO_WRITE;
|
||||
result.io_context.query_id = result.query_id.get();
|
||||
return result;
|
||||
}
|
||||
|
||||
FileRangeReadScheduler::FileRangeReadScheduler(FileRangeReadSchedulerOptions options,
|
||||
ThreadPool* executor)
|
||||
: _options(options),
|
||||
_global_budget(std::make_shared<FileRangeReadBudget>(options.max_bytes_per_be)),
|
||||
_mem_tracker(MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::CACHE,
|
||||
"FileRangeReadScheduler")),
|
||||
_executor(executor) {
|
||||
DORIS_CHECK(_executor != nullptr);
|
||||
read_ahead_bvars().resident_bytes << 0;
|
||||
}
|
||||
|
||||
FileRangeReadScheduler::~FileRangeReadScheduler() {
|
||||
shutdown();
|
||||
}
|
||||
|
||||
Status FileRangeReadSchedulerOptions::validate() const {
|
||||
if (max_bytes_per_query == 0) {
|
||||
return Status::InvalidArgument("file range read query byte limit must be positive");
|
||||
}
|
||||
if (max_bytes_per_be == 0) {
|
||||
return Status::InvalidArgument("file range read BE byte limit must be positive");
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status FileRangeReadScheduler::create(const FileRangeReadSchedulerOptions& options,
|
||||
ThreadPool* executor,
|
||||
std::unique_ptr<FileRangeReadScheduler>* output) {
|
||||
DORIS_CHECK(output != nullptr);
|
||||
DORIS_CHECK(executor != nullptr);
|
||||
RETURN_IF_ERROR(options.validate());
|
||||
output->reset(new FileRangeReadScheduler(options, executor));
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
std::shared_ptr<FileRangeReadContext> FileRangeReadScheduler::create_context() const {
|
||||
return std::make_shared<FileRangeReadContext>(_options.max_bytes_per_query);
|
||||
}
|
||||
|
||||
FileRangeReadSubmitResult FileRangeReadScheduler::try_submit(
|
||||
const std::vector<FileRange>& ranges, FileReaderSPtr reader,
|
||||
FileRangeReadIOContext io_context, const std::shared_ptr<FileRangeReadContext>& context) {
|
||||
const auto record_rejection = [&](FileRangeReadRejectReason reason) {
|
||||
read_ahead_bvars().rejected_batches << 1;
|
||||
if (io_context.statistics != nullptr) {
|
||||
COUNTER_UPDATE(&io_context.statistics->rejected_batches, 1);
|
||||
if (reason == FileRangeReadRejectReason::QUERY_BYTE_LIMIT) {
|
||||
COUNTER_UPDATE(&io_context.statistics->query_budget_rejected_batches, 1);
|
||||
} else if (reason == FileRangeReadRejectReason::GLOBAL_BYTE_LIMIT) {
|
||||
COUNTER_UPDATE(&io_context.statistics->be_budget_rejected_batches, 1);
|
||||
}
|
||||
}
|
||||
};
|
||||
const auto reject = [&](FileRangeReadRejectReason reason, Status status = Status::OK()) {
|
||||
record_rejection(reason);
|
||||
return FileRangeReadSubmitResult::rejected(reason, std::move(status));
|
||||
};
|
||||
Status validation_status = _validate_request(ranges, reader, context);
|
||||
if (!validation_status.ok()) {
|
||||
return reject(FileRangeReadRejectReason::INVALID_REQUEST, std::move(validation_status));
|
||||
}
|
||||
if (!accepting()) {
|
||||
return reject(FileRangeReadRejectReason::SHUTTING_DOWN,
|
||||
Status::Cancelled("file range read scheduler is shutting down"));
|
||||
}
|
||||
if (context->cancelled()) {
|
||||
return reject(FileRangeReadRejectReason::QUERY_CANCELLED,
|
||||
Status::Cancelled("query cancelled before file range submission"));
|
||||
}
|
||||
|
||||
std::vector<std::optional<FileRangeReadReservation>> reservations(ranges.size());
|
||||
const auto reserve_reason = _reserve_batch(context, ranges, &reservations);
|
||||
|
||||
std::vector<std::shared_ptr<FileRangeRead>> reads(ranges.size());
|
||||
std::vector<ReadTask> tasks;
|
||||
tasks.reserve(ranges.size());
|
||||
for (size_t index = 0; index < ranges.size(); ++index) {
|
||||
// Keep the input slot empty when admission skipped a range, so page mappings stay valid.
|
||||
if (!reservations[index].has_value()) {
|
||||
continue;
|
||||
}
|
||||
std::shared_ptr<FileRangeRead> read;
|
||||
Status status =
|
||||
FileRangeRead::create(ranges[index], _mem_tracker, std::move(*reservations[index]),
|
||||
&read, io_context.statistics);
|
||||
if (!status.ok()) {
|
||||
return reject(FileRangeReadRejectReason::ALLOCATION_FAILED, std::move(status));
|
||||
}
|
||||
tasks.push_back(ReadTask {
|
||||
.reader = reader, .read = read, .context = context, .io_context = io_context});
|
||||
reads[index] = std::move(read);
|
||||
}
|
||||
// Every range may exceed the remaining budgets, leaving the whole request on the sync path.
|
||||
if (tasks.empty()) {
|
||||
DORIS_CHECK(reserve_reason != FileRangeReadRejectReason::NONE);
|
||||
return reject(reserve_reason);
|
||||
}
|
||||
context->_register_reads(reads);
|
||||
|
||||
FileRangeReadRejectReason submit_reason = FileRangeReadRejectReason::NONE;
|
||||
Status submit_status = _submit_tasks(std::move(tasks), &submit_reason);
|
||||
if (!submit_status.ok()) {
|
||||
return reject(submit_reason, std::move(submit_status));
|
||||
}
|
||||
// Record a partially budget-rejected batch once, after execution admission succeeds. If every
|
||||
// executor task was rejected, the failure above already counted this submission.
|
||||
if (reserve_reason != FileRangeReadRejectReason::NONE) {
|
||||
record_rejection(reserve_reason);
|
||||
}
|
||||
return {.reads = std::move(reads),
|
||||
.reject_reason = FileRangeReadRejectReason::NONE,
|
||||
.status = Status::OK()};
|
||||
}
|
||||
|
||||
Status FileRangeReadScheduler::_validate_request(
|
||||
const std::vector<FileRange>& ranges, const FileReaderSPtr& reader,
|
||||
const std::shared_ptr<FileRangeReadContext>& context) const {
|
||||
if (reader == nullptr) {
|
||||
return Status::InvalidArgument("file range reader is null");
|
||||
}
|
||||
if (context == nullptr) {
|
||||
return Status::InvalidArgument("file range read context is null");
|
||||
}
|
||||
if (ranges.empty()) {
|
||||
return Status::InvalidArgument("file range read batch is empty");
|
||||
}
|
||||
|
||||
for (const auto& range : ranges) {
|
||||
if (range.size == 0) {
|
||||
return Status::InvalidArgument("file range at offset {} is empty", range.offset);
|
||||
}
|
||||
if (add_overflows(range.offset, range.size)) {
|
||||
return Status::InvalidArgument("file range at offset {} with size {} overflows",
|
||||
range.offset, range.size);
|
||||
}
|
||||
if (range.end() > reader->size()) {
|
||||
return Status::InvalidArgument("file range [{}, {}) exceeds file size {}", range.offset,
|
||||
range.end(), reader->size());
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
FileRangeReadRejectReason FileRangeReadScheduler::_reserve_batch(
|
||||
const std::shared_ptr<FileRangeReadContext>& context, const std::vector<FileRange>& ranges,
|
||||
std::vector<std::optional<FileRangeReadReservation>>* reservations) {
|
||||
DORIS_CHECK(reservations != nullptr);
|
||||
DORIS_CHECK(reservations->size() == ranges.size());
|
||||
auto& query_budget = context->_budget;
|
||||
std::scoped_lock lock(query_budget._mutex, _global_budget->_mutex);
|
||||
|
||||
size_t query_remaining = query_budget._max_bytes - query_budget._resident_bytes;
|
||||
size_t be_remaining = _global_budget->_max_bytes - _global_budget->_resident_bytes;
|
||||
size_t reserved_bytes = 0;
|
||||
auto reject_reason = FileRangeReadRejectReason::NONE;
|
||||
for (size_t index = 0; index < ranges.size(); ++index) {
|
||||
const size_t bytes = ranges[index].size;
|
||||
const auto reason = bytes > query_remaining ? FileRangeReadRejectReason::QUERY_BYTE_LIMIT
|
||||
: bytes > be_remaining ? FileRangeReadRejectReason::GLOBAL_BYTE_LIMIT
|
||||
: FileRangeReadRejectReason::NONE;
|
||||
// With 3 MiB available and ranges of 2, 2, and 1 MiB, admit the first and third ranges.
|
||||
// A range that cannot fit must not prevent later, smaller ranges from using the budget.
|
||||
if (reason != FileRangeReadRejectReason::NONE) {
|
||||
if (reject_reason == FileRangeReadRejectReason::NONE) {
|
||||
reject_reason = reason;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
(*reservations)[index].emplace(FileRangeReadReservation(context, _global_budget, bytes));
|
||||
query_remaining -= bytes;
|
||||
be_remaining -= bytes;
|
||||
reserved_bytes += bytes;
|
||||
}
|
||||
|
||||
query_budget._reserve_unlocked(reserved_bytes);
|
||||
_global_budget->_reserve_unlocked(reserved_bytes);
|
||||
read_ahead_bvars().resident_bytes << static_cast<int64_t>(reserved_bytes);
|
||||
return reject_reason;
|
||||
}
|
||||
|
||||
Status FileRangeReadScheduler::_submit_tasks(std::vector<ReadTask> tasks,
|
||||
FileRangeReadRejectReason* reject_reason) {
|
||||
DORIS_CHECK(!tasks.empty());
|
||||
DORIS_CHECK(reject_reason != nullptr);
|
||||
std::lock_guard lock(_mutex);
|
||||
if (!_accepting) {
|
||||
*reject_reason = FileRangeReadRejectReason::SHUTTING_DOWN;
|
||||
return Status::Cancelled("file range read scheduler is shutting down");
|
||||
}
|
||||
|
||||
size_t submitted_tasks = 0;
|
||||
Status first_failure;
|
||||
for (auto& task : tasks) {
|
||||
auto read = task.read;
|
||||
DORIS_CHECK(_inflight_reads.emplace(read).second);
|
||||
read_ahead_bvars().inflight_ranges << 1;
|
||||
task.submit_ns = MonotonicNanos();
|
||||
Status status = _executor->submit_func([this, task = std::move(task), read]() mutable {
|
||||
Defer task_finished {[this, read]() { _task_finished(read); }};
|
||||
_run_task(std::move(task));
|
||||
});
|
||||
if (status.ok()) {
|
||||
++submitted_tasks;
|
||||
if (read->_statistics != nullptr) {
|
||||
COUNTER_UPDATE(&read->_statistics->submitted_ranges, 1);
|
||||
COUNTER_UPDATE(&read->_statistics->submitted_bytes,
|
||||
static_cast<int64_t>(read->range().size));
|
||||
}
|
||||
continue;
|
||||
}
|
||||
DORIS_CHECK(_inflight_reads.erase(read) == 1);
|
||||
read_ahead_bvars().inflight_ranges << -1;
|
||||
if (first_failure.ok()) {
|
||||
first_failure = status;
|
||||
}
|
||||
if (read->_statistics != nullptr) {
|
||||
COUNTER_UPDATE(&read->_statistics->submit_failed_ranges, 1);
|
||||
}
|
||||
read->_publish_submit_failure(std::move(status));
|
||||
}
|
||||
|
||||
if (submitted_tasks == 0) {
|
||||
DORIS_CHECK(!first_failure.ok());
|
||||
*reject_reason = FileRangeReadRejectReason::EXECUTOR_REJECTED;
|
||||
return first_failure;
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
void FileRangeReadScheduler::_task_finished(const std::shared_ptr<FileRangeRead>& read) {
|
||||
std::lock_guard lock(_mutex);
|
||||
DORIS_CHECK(_inflight_reads.erase(read) == 1);
|
||||
read_ahead_bvars().inflight_ranges << -1;
|
||||
if (_inflight_reads.empty()) {
|
||||
_tasks_cv.notify_all();
|
||||
}
|
||||
}
|
||||
|
||||
void FileRangeReadScheduler::_run_task(ReadTask task) const {
|
||||
if (!task.read->_mark_running()) {
|
||||
// Count queued cancellation only for tasks accepted by the executor. A handle may be
|
||||
// cancelled before submission, including a submission the executor later rejects.
|
||||
if (task.read->_statistics != nullptr) {
|
||||
COUNTER_UPDATE(&task.read->_statistics->cancelled_ranges, 1);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
FileRangeReadStats stats;
|
||||
stats.queue_wait_ns = MonotonicNanos() - task.submit_ns;
|
||||
if (task.context->cancelled() || !accepting()) {
|
||||
task.read->_publish_cancelled(std::move(stats));
|
||||
return;
|
||||
}
|
||||
task.io_context.io_context.file_cache_stats = &stats.file_cache;
|
||||
task.io_context.io_context.file_reader_stats = &stats.file_reader;
|
||||
Status status;
|
||||
const int64_t read_start_ns = MonotonicNanos();
|
||||
try {
|
||||
status = task.reader->read_at(task.read->range().offset,
|
||||
Slice(task.read->_data, task.read->range().size),
|
||||
&stats.bytes_read, &task.io_context.io_context);
|
||||
} catch (const Exception& exception) {
|
||||
status = exception.to_status();
|
||||
} catch (const std::exception& exception) {
|
||||
status = Status::IOError("file range read threw an exception: {}", exception.what());
|
||||
} catch (...) {
|
||||
status = Status::IOError("file range read threw an unknown exception");
|
||||
}
|
||||
stats.read_ns = MonotonicNanos() - read_start_ns;
|
||||
if (status.ok() && stats.bytes_read == task.read->range().size) {
|
||||
stats.successful_source_bytes = stats.bytes_read;
|
||||
}
|
||||
|
||||
if (task.context->cancelled() || task.read->_is_cancel_requested() || !accepting()) {
|
||||
task.read->_publish_cancelled(std::move(stats));
|
||||
} else if (!status.ok()) {
|
||||
task.read->_publish_failed(std::move(status), std::move(stats));
|
||||
} else if (stats.bytes_read != task.read->range().size) {
|
||||
const size_t bytes_read = stats.bytes_read;
|
||||
task.read->_publish_failed(Status::IOError("file range read returned {} of {} bytes",
|
||||
bytes_read, task.read->range().size),
|
||||
std::move(stats));
|
||||
} else {
|
||||
TEST_SYNC_POINT("FileRangeReadScheduler::_run_task:before_publish_ready");
|
||||
task.read->_publish_ready(std::move(stats));
|
||||
}
|
||||
}
|
||||
|
||||
void FileRangeReadScheduler::shutdown() {
|
||||
std::unique_lock lock(_mutex);
|
||||
_accepting = false;
|
||||
for (const auto& read : _inflight_reads) {
|
||||
read->request_cancel();
|
||||
}
|
||||
_tasks_cv.wait(lock, [this]() { return _inflight_reads.empty(); });
|
||||
}
|
||||
|
||||
bool FileRangeReadScheduler::accepting() const {
|
||||
std::lock_guard lock(_mutex);
|
||||
return _accepting;
|
||||
}
|
||||
|
||||
} // namespace doris::io
|
||||
@@ -0,0 +1,306 @@
|
||||
// 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <atomic>
|
||||
#include <condition_variable>
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <optional>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "common/status.h"
|
||||
#include "io/fs/file_range_coalescer.h"
|
||||
#include "io/fs/file_reader.h"
|
||||
#include "io/io_common.h"
|
||||
#include "util/slice.h"
|
||||
|
||||
namespace doris {
|
||||
|
||||
class MemTrackerLimiter;
|
||||
class ThreadPool;
|
||||
|
||||
namespace io {
|
||||
|
||||
struct ReadAheadStatistics;
|
||||
|
||||
enum class FileRangeReadRejectReason : uint8_t {
|
||||
/// At least one executor task was accepted. Other ranges may be skipped or fail to submit/read.
|
||||
NONE,
|
||||
/// A reader, context, or range failed request validation.
|
||||
INVALID_REQUEST,
|
||||
/// The query context had already been cancelled before admission.
|
||||
QUERY_CANCELLED,
|
||||
/// The query-scoped byte budget prevented range admission.
|
||||
QUERY_BYTE_LIMIT,
|
||||
/// The BE-wide byte budget prevented range admission.
|
||||
GLOBAL_BYTE_LIMIT,
|
||||
/// At least one range buffer could not be allocated; no read handles are returned.
|
||||
ALLOCATION_FAILED,
|
||||
/// The executor rejected every task in the batch.
|
||||
EXECUTOR_REJECTED,
|
||||
/// The scheduler stopped accepting work before the batch was submitted.
|
||||
SHUTTING_DOWN,
|
||||
};
|
||||
|
||||
struct FileRangeReadSchedulerOptions {
|
||||
/// Maximum bytes retained by live range handles from one query.
|
||||
size_t max_bytes_per_query {1};
|
||||
/// Maximum bytes retained by all live range handles in this BE.
|
||||
size_t max_bytes_per_be {1};
|
||||
|
||||
Status validate() const;
|
||||
};
|
||||
|
||||
/// Byte reservation shared by range buffers. It limits retained memory, not executor concurrency.
|
||||
class FileRangeReadBudget {
|
||||
public:
|
||||
explicit FileRangeReadBudget(size_t max_bytes);
|
||||
|
||||
size_t resident_bytes() const;
|
||||
|
||||
private:
|
||||
friend class FileRangeReadScheduler;
|
||||
friend class FileRangeReadReservation;
|
||||
|
||||
bool _can_reserve_unlocked(size_t bytes) const;
|
||||
void _reserve_unlocked(size_t bytes);
|
||||
void _release(size_t bytes);
|
||||
|
||||
mutable std::mutex _mutex;
|
||||
const size_t _max_bytes;
|
||||
size_t _resident_bytes {0};
|
||||
};
|
||||
|
||||
class FileRangeRead;
|
||||
|
||||
/// Query-scoped cancellation and admission state shared by every segment read.
|
||||
class FileRangeReadContext {
|
||||
public:
|
||||
explicit FileRangeReadContext(size_t max_bytes);
|
||||
|
||||
/// Cancel all registered reads. Running source IO finishes normally and publishes CANCELLED.
|
||||
void cancel();
|
||||
bool cancelled() const { return _cancelled.load(std::memory_order_acquire); }
|
||||
size_t resident_bytes() const { return _budget.resident_bytes(); }
|
||||
|
||||
private:
|
||||
friend class FileRangeReadScheduler;
|
||||
friend class FileRangeReadReservation;
|
||||
|
||||
void _register_reads(const std::vector<std::shared_ptr<FileRangeRead>>& reads);
|
||||
|
||||
FileRangeReadBudget _budget;
|
||||
std::atomic<bool> _cancelled {false};
|
||||
mutable std::mutex _reads_mutex;
|
||||
std::vector<std::weak_ptr<FileRangeRead>> _reads;
|
||||
};
|
||||
|
||||
/// Move-only reservation that releases both query and BE byte budgets on destruction.
|
||||
class FileRangeReadReservation {
|
||||
public:
|
||||
FileRangeReadReservation(const FileRangeReadReservation&) = delete;
|
||||
FileRangeReadReservation& operator=(const FileRangeReadReservation&) = delete;
|
||||
FileRangeReadReservation(FileRangeReadReservation&& other) noexcept;
|
||||
FileRangeReadReservation& operator=(FileRangeReadReservation&& other) noexcept;
|
||||
~FileRangeReadReservation();
|
||||
|
||||
private:
|
||||
friend class FileRangeReadScheduler;
|
||||
|
||||
FileRangeReadReservation(std::shared_ptr<FileRangeReadContext> context,
|
||||
std::shared_ptr<FileRangeReadBudget> global_budget,
|
||||
size_t bytes) noexcept;
|
||||
void _reset();
|
||||
|
||||
std::shared_ptr<FileRangeReadContext> _context;
|
||||
std::shared_ptr<FileRangeReadBudget> _global_budget;
|
||||
size_t _bytes {0};
|
||||
};
|
||||
|
||||
struct FileRangeReadStats {
|
||||
size_t bytes_read {0};
|
||||
/// Full successful source reads, including a read whose buffer is subsequently cancelled.
|
||||
size_t successful_source_bytes {0};
|
||||
FileCacheStatistics file_cache;
|
||||
FileReaderStats file_reader;
|
||||
int64_t queue_wait_ns {0};
|
||||
int64_t read_ns {0};
|
||||
};
|
||||
|
||||
/// One scheduled range and its owned buffer. The byte reservation follows this handle's lifetime.
|
||||
class FileRangeRead {
|
||||
public:
|
||||
enum class State : uint8_t {
|
||||
/// Admitted and allocated, but its executor task has not started.
|
||||
QUEUED,
|
||||
/// The executor task owns the handle and may be blocked in source IO.
|
||||
RUNNING,
|
||||
/// The entire requested range is available through data() or slice().
|
||||
READY,
|
||||
/// Submission, source IO, or exact-length validation failed.
|
||||
FAILED,
|
||||
/// Cancellation won before execution or was observed after source IO returned.
|
||||
CANCELLED,
|
||||
};
|
||||
|
||||
~FileRangeRead();
|
||||
|
||||
/// Wait for a terminal state and return its status. Optionally report whether the buffer was
|
||||
/// already READY on entry, using the same lock as the wait.
|
||||
Status wait(bool* ready_on_entry = nullptr);
|
||||
/// Request cancellation without interrupting an already running source read.
|
||||
void request_cancel();
|
||||
State state() const;
|
||||
const FileRange& range() const { return _range; }
|
||||
/// Return the full buffer. The handle must be READY and owns the returned memory.
|
||||
Slice data() const;
|
||||
/// Return a checked subrange of a READY buffer.
|
||||
Slice slice(size_t buffer_offset, size_t size) const;
|
||||
FileRangeReadStats stats() const;
|
||||
|
||||
private:
|
||||
friend class FileRangeReadScheduler;
|
||||
|
||||
FileRangeRead(FileRange range, std::shared_ptr<MemTrackerLimiter> tracker,
|
||||
FileRangeReadReservation reservation,
|
||||
std::shared_ptr<ReadAheadStatistics> statistics);
|
||||
static Status create(FileRange range, std::shared_ptr<MemTrackerLimiter> tracker,
|
||||
FileRangeReadReservation reservation,
|
||||
std::shared_ptr<FileRangeRead>* output,
|
||||
std::shared_ptr<ReadAheadStatistics> statistics = nullptr);
|
||||
|
||||
static bool _is_terminal(State state);
|
||||
bool _mark_running();
|
||||
bool _is_cancel_requested() const;
|
||||
void _publish_ready(FileRangeReadStats stats);
|
||||
void _publish_failed(Status status, FileRangeReadStats stats);
|
||||
void _publish_submit_failure(Status status);
|
||||
void _publish_cancelled(FileRangeReadStats stats = {});
|
||||
void _publish_from_running(State state, Status status, FileRangeReadStats stats);
|
||||
const FileRange _range;
|
||||
char* _data {nullptr};
|
||||
const std::shared_ptr<MemTrackerLimiter> _tracker;
|
||||
FileRangeReadReservation _reservation;
|
||||
const std::shared_ptr<ReadAheadStatistics> _statistics;
|
||||
mutable std::mutex _mutex;
|
||||
std::condition_variable _cv;
|
||||
State _state {State::QUEUED};
|
||||
Status _status;
|
||||
FileRangeReadStats _stats;
|
||||
bool _cancel_requested {false};
|
||||
};
|
||||
|
||||
/// Owns every pointer embedded in IOContext that an asynchronous read keeps after submission.
|
||||
struct FileRangeReadIOContext {
|
||||
IOContext io_context;
|
||||
std::shared_ptr<TUniqueId> query_id;
|
||||
std::shared_ptr<ReadAheadStatistics> statistics;
|
||||
|
||||
/// Copy fields safe for asynchronous use, own query_id, and detach caller-owned statistics.
|
||||
static FileRangeReadIOContext from_caller(const IOContext& source);
|
||||
};
|
||||
|
||||
struct FileRangeReadSubmitResult {
|
||||
/// One slot per input range on acceptance; a null slot was skipped by byte-budget admission.
|
||||
/// A completely rejected submission returns an empty vector.
|
||||
std::vector<std::shared_ptr<FileRangeRead>> reads;
|
||||
FileRangeReadRejectReason reject_reason {FileRangeReadRejectReason::NONE};
|
||||
Status status;
|
||||
|
||||
static FileRangeReadSubmitResult rejected(FileRangeReadRejectReason reason,
|
||||
Status status = Status::OK());
|
||||
bool accepted() const {
|
||||
return reject_reason == FileRangeReadRejectReason::NONE && status.ok();
|
||||
}
|
||||
};
|
||||
|
||||
/// BE-global admission and execution scheduler for already-planned file ranges. It does not
|
||||
/// understand pages, range coalescing, cache-block completion, or writeback policy.
|
||||
class FileRangeReadScheduler {
|
||||
public:
|
||||
~FileRangeReadScheduler();
|
||||
|
||||
FileRangeReadScheduler(const FileRangeReadScheduler&) = delete;
|
||||
FileRangeReadScheduler& operator=(const FileRangeReadScheduler&) = delete;
|
||||
|
||||
/// `executor` must remain running and outlive the scheduler. Its thread count controls range
|
||||
/// read concurrency; the scheduler never shuts down the shared executor.
|
||||
static Status create(const FileRangeReadSchedulerOptions& options, ThreadPool* executor,
|
||||
std::unique_ptr<FileRangeReadScheduler>* output);
|
||||
|
||||
/// Create a query-scoped cancellation and byte-budget context.
|
||||
std::shared_ptr<FileRangeReadContext> create_context() const;
|
||||
|
||||
/// Admit ranges in input order under both byte-budget locks, skipping ranges that do not fit,
|
||||
/// then allocate and submit one task per admitted range. Returned slots preserve input indices;
|
||||
/// unadmitted ranges have null handles. Allocation failure rolls back all admitted buffers.
|
||||
/// Executor rejection is exposed as a failed handle when other tasks were accepted.
|
||||
FileRangeReadSubmitResult try_submit(const std::vector<FileRange>& ranges,
|
||||
FileReaderSPtr reader, FileRangeReadIOContext io_context,
|
||||
const std::shared_ptr<FileRangeReadContext>& context);
|
||||
|
||||
/// Stop admission, request cancellation, and wait for all accepted executor tasks. Idempotent.
|
||||
void shutdown();
|
||||
bool accepting() const;
|
||||
std::shared_ptr<FileRangeReadBudget> global_budget() const { return _global_budget; }
|
||||
|
||||
private:
|
||||
struct ReadTask {
|
||||
FileReaderSPtr reader;
|
||||
std::shared_ptr<FileRangeRead> read;
|
||||
std::shared_ptr<FileRangeReadContext> context;
|
||||
FileRangeReadIOContext io_context;
|
||||
int64_t submit_ns {0};
|
||||
};
|
||||
|
||||
FileRangeReadScheduler(FileRangeReadSchedulerOptions options, ThreadPool* executor);
|
||||
|
||||
/// Validate the reader, context, and individual ranges before reserving any memory.
|
||||
Status _validate_request(const std::vector<FileRange>& ranges, const FileReaderSPtr& reader,
|
||||
const std::shared_ptr<FileRangeReadContext>& context) const;
|
||||
/// Reserve query and BE bytes for ranges that fit, preserving input indices. Return the first
|
||||
/// encountered byte-budget rejection, including when some ranges were admitted.
|
||||
FileRangeReadRejectReason _reserve_batch(
|
||||
const std::shared_ptr<FileRangeReadContext>& context,
|
||||
const std::vector<FileRange>& ranges,
|
||||
std::vector<std::optional<FileRangeReadReservation>>* reservations);
|
||||
/// Submit one task per range while retaining accepted handles for shutdown coordination.
|
||||
Status _submit_tasks(std::vector<ReadTask> tasks, FileRangeReadRejectReason* reject_reason);
|
||||
void _task_finished(const std::shared_ptr<FileRangeRead>& read);
|
||||
/// Perform one exact source read and publish one terminal handle state.
|
||||
void _run_task(ReadTask task) const;
|
||||
|
||||
const FileRangeReadSchedulerOptions _options;
|
||||
const std::shared_ptr<FileRangeReadBudget> _global_budget;
|
||||
const std::shared_ptr<MemTrackerLimiter> _mem_tracker;
|
||||
ThreadPool* const _executor;
|
||||
|
||||
mutable std::mutex _mutex;
|
||||
std::condition_variable _tasks_cv;
|
||||
// The shared executor owns task ordering. This set only keeps accepted reads alive and allows
|
||||
// shutdown to cancel them; its size is not a concurrency limit.
|
||||
std::unordered_set<std::shared_ptr<FileRangeRead>> _inflight_reads;
|
||||
bool _accepting {true};
|
||||
};
|
||||
|
||||
} // namespace io
|
||||
} // namespace doris
|
||||
@@ -55,6 +55,8 @@ struct ReadAheadStatistics {
|
||||
RuntimeProfile::Counter planned_bytes {TUnit::BYTES};
|
||||
RuntimeProfile::Counter submitted_ranges {TUnit::UNIT};
|
||||
RuntimeProfile::Counter submitted_bytes {TUnit::BYTES};
|
||||
/// Whole-batch rejections and byte-budget partial rejections. The budget counters below also
|
||||
/// include partially admitted batches; their cause is the first budget rejection in that batch.
|
||||
RuntimeProfile::Counter rejected_batches {TUnit::UNIT};
|
||||
RuntimeProfile::Counter query_budget_rejected_batches {TUnit::UNIT};
|
||||
RuntimeProfile::Counter be_budget_rejected_batches {TUnit::UNIT};
|
||||
|
||||
@@ -159,7 +159,7 @@ Status S3FileReader::read_at_impl(size_t offset, Slice result, size_t* bytes_rea
|
||||
.count();
|
||||
s3_file_reader_latency << (end_ts - begin_ts);
|
||||
}};
|
||||
SCOPED_RAW_TIMER(&_s3_stats.total_get_request_time_ns);
|
||||
SCOPED_ATOMIC_TIMER(&_s3_stats.total_get_request_time_ns);
|
||||
|
||||
int total_sleep_time = 0;
|
||||
Status last_error;
|
||||
|
||||
@@ -63,11 +63,12 @@ protected:
|
||||
|
||||
private:
|
||||
struct S3Statistics {
|
||||
int64_t total_get_request_counter = 0;
|
||||
int64_t too_many_request_err_counter = 0;
|
||||
int64_t too_many_request_sleep_time_ms = 0;
|
||||
int64_t total_bytes_read = 0;
|
||||
int64_t total_get_request_time_ns = 0;
|
||||
// The reader is shared by foreground reads and parallel background hole-fill GETs.
|
||||
std::atomic<int64_t> total_get_request_counter {0};
|
||||
std::atomic<int64_t> too_many_request_err_counter {0};
|
||||
std::atomic<int64_t> too_many_request_sleep_time_ms {0};
|
||||
std::atomic<int64_t> total_bytes_read {0};
|
||||
std::atomic<int64_t> total_get_request_time_ns {0};
|
||||
};
|
||||
Path _path;
|
||||
size_t _file_size;
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,146 @@
|
||||
// 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 "io/fs/s3_file_reader.h"
|
||||
|
||||
#include <aws/s3/S3Client.h>
|
||||
#include <aws/s3/model/GetObjectRequest.h>
|
||||
#include <aws/s3/model/GetObjectResult.h>
|
||||
#include <gmock/gmock.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <array>
|
||||
#include <chrono>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
|
||||
#include "common/config.h"
|
||||
#include "cpp/obj-client/s3_obj_storage_client.h"
|
||||
#include "runtime/runtime_profile.h"
|
||||
#include "util/countdown_latch.h"
|
||||
#include "util/defer_op.h"
|
||||
#include "util/s3_util.h"
|
||||
#include "util/threadpool.h"
|
||||
|
||||
namespace doris::io {
|
||||
|
||||
class MockRangeReadS3Client final : public Aws::S3::S3Client {
|
||||
public:
|
||||
MOCK_METHOD(Aws::S3::Model::GetObjectOutcome, GetObject,
|
||||
(const Aws::S3::Model::GetObjectRequest&), (const, override));
|
||||
};
|
||||
|
||||
namespace {
|
||||
|
||||
using namespace std::chrono_literals;
|
||||
|
||||
Aws::S3::Model::GetObjectOutcome successful_read(const Aws::S3::Model::GetObjectRequest& request) {
|
||||
auto* stream = request.GetResponseStreamFactory()();
|
||||
*stream << "0123456789abcdef";
|
||||
Aws::S3::Model::GetObjectResult result;
|
||||
result.ReplaceBody(stream);
|
||||
result.SetContentLength(16);
|
||||
return Aws::S3::Model::GetObjectOutcome(std::move(result));
|
||||
}
|
||||
|
||||
class S3FileReaderTest : public testing::Test {
|
||||
protected:
|
||||
static void SetUpTestSuite() { S3ClientFactory::instance(); }
|
||||
};
|
||||
|
||||
TEST_F(S3FileReaderTest, ConcurrentReadsCollectAllStatistics) {
|
||||
constexpr size_t thread_count = 8;
|
||||
constexpr size_t reads_per_thread = 100;
|
||||
auto client = std::make_shared<MockRangeReadS3Client>();
|
||||
EXPECT_CALL(*client, GetObject(testing::_))
|
||||
.Times(thread_count * reads_per_thread)
|
||||
.WillRepeatedly(successful_read);
|
||||
auto holder = std::make_shared<ObjClientHolder>(S3ClientConf {});
|
||||
holder->_client = std::make_shared<S3ObjStorageClient>(client);
|
||||
RuntimeProfile profile("parallel-s3-reads");
|
||||
S3FileReader reader(holder, "bucket", "key", 16 * thread_count, &profile);
|
||||
std::unique_ptr<ThreadPool> pool;
|
||||
ASSERT_TRUE(
|
||||
ThreadPoolBuilder("S3FileReaderTest").set_max_threads(thread_count).build(&pool).ok());
|
||||
CountDownLatch ready(thread_count);
|
||||
std::array<Status, thread_count> results;
|
||||
for (size_t index = 0; index < thread_count; ++index) {
|
||||
ASSERT_TRUE(pool->submit_func([&, index]() {
|
||||
ready.count_down();
|
||||
EXPECT_TRUE(ready.wait_for(5s));
|
||||
for (size_t iteration = 0; iteration < reads_per_thread; ++iteration) {
|
||||
std::string buffer(16, '\0');
|
||||
size_t bytes_read = 0;
|
||||
results[index] =
|
||||
reader.read_at(index * 16, Slice(buffer), &bytes_read);
|
||||
if (!results[index].ok()) {
|
||||
return;
|
||||
}
|
||||
EXPECT_EQ(bytes_read, 16);
|
||||
EXPECT_EQ(buffer, "0123456789abcdef");
|
||||
}
|
||||
}).ok());
|
||||
}
|
||||
pool->wait();
|
||||
for (const auto& result : results) {
|
||||
ASSERT_TRUE(result.ok()) << result;
|
||||
}
|
||||
reader._collect_profile_before_close();
|
||||
EXPECT_EQ(profile.get_counter("TotalGetRequest")->value(), thread_count * reads_per_thread);
|
||||
EXPECT_EQ(profile.get_counter("TotalBytesRead")->value(), 16 * thread_count * reads_per_thread);
|
||||
EXPECT_GT(profile.get_counter("TotalGetRequestTime")->value(), 0);
|
||||
EXPECT_EQ(profile.get_counter("TooManyRequestErr")->value(), 0);
|
||||
}
|
||||
|
||||
TEST_F(S3FileReaderTest, CollectsRetryStatistics) {
|
||||
const auto old_retries = config::max_s3_client_retry;
|
||||
const auto old_base_wait = config::s3_read_base_wait_time_ms;
|
||||
const auto old_max_wait = config::s3_read_max_wait_time_ms;
|
||||
Defer restore {[&]() {
|
||||
config::max_s3_client_retry = old_retries;
|
||||
config::s3_read_base_wait_time_ms = old_base_wait;
|
||||
config::s3_read_max_wait_time_ms = old_max_wait;
|
||||
}};
|
||||
config::max_s3_client_retry = 1;
|
||||
config::s3_read_base_wait_time_ms = 1;
|
||||
config::s3_read_max_wait_time_ms = 2;
|
||||
auto client = std::make_shared<MockRangeReadS3Client>();
|
||||
EXPECT_CALL(*client, GetObject(testing::_))
|
||||
.WillOnce([](const Aws::S3::Model::GetObjectRequest&) {
|
||||
Aws::S3::S3Error error;
|
||||
error.SetResponseCode(Aws::Http::HttpResponseCode::TOO_MANY_REQUESTS);
|
||||
return Aws::S3::Model::GetObjectOutcome(std::move(error));
|
||||
})
|
||||
.WillOnce(successful_read);
|
||||
auto holder = std::make_shared<ObjClientHolder>(S3ClientConf {});
|
||||
holder->_client = std::make_shared<S3ObjStorageClient>(client);
|
||||
RuntimeProfile profile("retry-s3-read");
|
||||
S3FileReader reader(holder, "bucket", "key", 16, &profile);
|
||||
std::string buffer(16, '\0');
|
||||
size_t bytes_read = 0;
|
||||
ASSERT_TRUE(reader.read_at(0, Slice(buffer), &bytes_read).ok());
|
||||
EXPECT_EQ(buffer, "0123456789abcdef");
|
||||
reader._collect_profile_before_close();
|
||||
EXPECT_EQ(profile.get_counter("TotalGetRequest")->value(), 2);
|
||||
EXPECT_EQ(profile.get_counter("TotalBytesRead")->value(), 16);
|
||||
EXPECT_EQ(profile.get_counter("TooManyRequestErr")->value(), 1);
|
||||
EXPECT_EQ(profile.get_counter("TooManyRequestSleepTime")->value(), 2);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace doris::io
|
||||
Reference in new issue
Block a user