diff --git a/be/src/exec/sink/writer/varrow_flight_result_writer.cpp b/be/src/exec/sink/writer/varrow_flight_result_writer.cpp index f4729f90ce0..867a8608c2b 100644 --- a/be/src/exec/sink/writer/varrow_flight_result_writer.cpp +++ b/be/src/exec/sink/writer/varrow_flight_result_writer.cpp @@ -21,8 +21,11 @@ #include #include +#include + #include "core/block/block.h" #include "exprs/vexpr_context.h" +#include "runtime/query_context.h" #include "runtime/result_block_buffer.h" #include "runtime/runtime_state.h" #include "runtime/thread_context.h" @@ -96,15 +99,44 @@ Status ArrowFlightResultBlockBuffer::get_schema(std::shared_ptr* print_id(_fragment_id), _status)); } -Status ArrowFlightResultBlockBuffer::get_arrow_batch(std::shared_ptr* result) { +ArrowFlightResultBlockBuffer::ArrowFlightResultBlockBuffer(TUniqueId id, RuntimeState* state, + std::shared_ptr schema, + int buffer_size) + : ResultBlockBuffer(id, state, buffer_size), + _arrow_schema(std::move(schema)), + _profile("ResultBlockBuffer " + print_id(_fragment_id)), + _timezone_obj(state->timezone_obj()), + _query_id(state->query_id()), + _query_ctx(state->get_query_ctx() ? state->get_query_ctx()->weak_from_this() + : std::weak_ptr {}) { + _serialize_batch_ns_timer = ADD_TIMER(&_profile, "SerializeBatchNsTime"); + _uncompressed_bytes_counter = ADD_COUNTER(&_profile, "UncompressedBytes", TUnit::BYTES); + _compressed_bytes_counter = ADD_COUNTER(&_profile, "CompressedBytes", TUnit::BYTES); +} + +void ArrowFlightResultBlockBuffer::cancel_query(const Status& reason) { + cancel(reason); + // The result buffer can be keyed by a fragment instance id rather than the query id. + // Keep only a weak query reference, and cancel outside the buffer/map locks. + if (auto query_ctx = _query_ctx.lock()) { + query_ctx->cancel(reason); + } +} + +Status ArrowFlightResultBlockBuffer::get_arrow_batch(std::shared_ptr* result, bool* eos) { + *result = nullptr; + *eos = false; std::unique_lock l(_lock); Defer defer {[&]() { _update_dependency(); }}; if (!_status.ok()) { return _status; } - while (_result_batch_queue.empty() && _status.ok() && !_is_close) { - _arrow_data_arrival.wait_for(l, std::chrono::milliseconds(20)); + if (!_arrow_data_arrival.wait_for(l, std::chrono::milliseconds(20), [&] { + return !_result_batch_queue.empty() || !_status.ok() || _is_close; + })) { + // Let the reader check RPC cancellation without treating an empty wait as EOF. + return Status::OK(); } if (!_status.ok()) { @@ -125,6 +157,7 @@ Status ArrowFlightResultBlockBuffer::get_arrow_batch(std::shared_ptr* res // normal path end if (_is_close) { + *eos = true; if (!_status.ok()) { return _status; } diff --git a/be/src/exec/sink/writer/varrow_flight_result_writer.h b/be/src/exec/sink/writer/varrow_flight_result_writer.h index 24b85a6cf7f..c348f1f63d7 100644 --- a/be/src/exec/sink/writer/varrow_flight_result_writer.h +++ b/be/src/exec/sink/writer/varrow_flight_result_writer.h @@ -26,6 +26,7 @@ namespace doris { class RuntimeState; class PFetchArrowDataResult; +class QueryContext; class Block; @@ -59,17 +60,12 @@ class ArrowFlightResultBlockBuffer final : public ResultBlockBuffer schema, int buffer_size) - : ResultBlockBuffer(id, state, buffer_size), - _arrow_schema(schema), - _profile("ResultBlockBuffer " + print_id(_fragment_id)), - _timezone_obj(state->timezone_obj()) { - _serialize_batch_ns_timer = ADD_TIMER(&_profile, "SerializeBatchNsTime"); - _uncompressed_bytes_counter = ADD_COUNTER(&_profile, "UncompressedBytes", TUnit::BYTES); - _compressed_bytes_counter = ADD_COUNTER(&_profile, "CompressedBytes", TUnit::BYTES); - } + std::shared_ptr schema, int buffer_size); ~ArrowFlightResultBlockBuffer() override = default; - Status get_arrow_batch(std::shared_ptr* result); + // A timed wait may return OK with no block and eos=false; the caller must retry. + Status get_arrow_batch(std::shared_ptr* result, bool* eos); + void cancel_query(const Status& reason); + const TUniqueId& query_id() const { return _query_id; } void get_timezone(cctz::time_zone& timezone_obj) { timezone_obj = _timezone_obj; } Status get_schema(std::shared_ptr* arrow_schema); @@ -82,6 +78,8 @@ private: RuntimeProfile::Counter* _uncompressed_bytes_counter = nullptr; RuntimeProfile::Counter* _compressed_bytes_counter = nullptr; cctz::time_zone _timezone_obj; + const TUniqueId _query_id; + std::weak_ptr _query_ctx; }; class VArrowFlightResultWriter final : public ResultWriter { diff --git a/be/src/runtime/result_buffer_mgr.cpp b/be/src/runtime/result_buffer_mgr.cpp index 0e53128c0b3..177b806c15a 100644 --- a/be/src/runtime/result_buffer_mgr.cpp +++ b/be/src/runtime/result_buffer_mgr.cpp @@ -93,6 +93,9 @@ Status ResultBufferMgr::create_sender(const TUniqueId& unique_id, int buffer_siz { std::unique_lock wlock(_buffer_map_lock); _buffer_map.insert(std::make_pair(unique_id, control_block)); + if (arrow_flight) { + _arrow_flight_query_buffers[state->query_id()].insert(unique_id); + } // ResultBlockBufferBase should destroy after max_timeout // for exceed max_timeout FE will return timeout to client // otherwise in some case may block all fragment handle threads @@ -135,12 +138,52 @@ bool ResultBufferMgr::cancel(const TUniqueId& unique_id, const Status& reason) { auto exist = _buffer_map.end() != iter; if (exist) { + if (auto arrow_buffer = + std::dynamic_pointer_cast(iter->second)) { + // create_sender and cancel_arrow_flight_buffers update both maps under this lock. + auto query = _arrow_flight_query_buffers.find(arrow_buffer->query_id()); + DORIS_CHECK(query != _arrow_flight_query_buffers.end()); + query->second.erase(unique_id); + if (query->second.empty()) { + _arrow_flight_query_buffers.erase(query); + } + } iter->second->cancel(reason); _buffer_map.erase(iter); } return exist; } +void ResultBufferMgr::cancel_arrow_flight_query(const TUniqueId& buffer_id, const Status& reason) { + std::shared_ptr buffer; + if (find_buffer(buffer_id, buffer).ok()) { + cancel_arrow_flight_buffers(buffer->query_id(), reason); + } +} + +void ResultBufferMgr::cancel_arrow_flight_buffers(const TUniqueId& query_id, const Status& reason) { + std::vector> buffers; + { + std::unique_lock lock(_buffer_map_lock); + auto query = _arrow_flight_query_buffers.find(query_id); + if (query == _arrow_flight_query_buffers.end()) { + return; + } + for (const auto& buffer_id : query->second) { + auto it = _buffer_map.find(buffer_id); + DORIS_CHECK(it != _buffer_map.end()); + buffers.push_back(std::static_pointer_cast(it->second)); + _buffer_map.erase(it); + } + _arrow_flight_query_buffers.erase(query); + } + // Clear all sibling buffers even if their local query context has already expired. + // Never cancel a query while holding the result-buffer map lock. + for (const auto& buffer : buffers) { + buffer->cancel_query(reason); + } +} + void ResultBufferMgr::cancel_at_time(time_t cancel_time, const TUniqueId& unique_id) { std::lock_guard l(_timeout_lock); auto iter = _timeout_map.find(cancel_time); diff --git a/be/src/runtime/result_buffer_mgr.h b/be/src/runtime/result_buffer_mgr.h index 96cb6a5b7e6..e0170c88176 100644 --- a/be/src/runtime/result_buffer_mgr.h +++ b/be/src/runtime/result_buffer_mgr.h @@ -27,6 +27,7 @@ #include #include #include +#include #include #include "common/status.h" @@ -71,6 +72,9 @@ public: // cancel bool cancel(const TUniqueId& unique_id, const Status& reason); + void cancel_arrow_flight_query(const TUniqueId& buffer_id, const Status& reason); + void cancel_arrow_flight_buffers(const TUniqueId& query_id, const Status& reason); + // cancel one query at a future time. void cancel_at_time(time_t cancel_time, const TUniqueId& unique_id); @@ -89,6 +93,7 @@ private: std::shared_mutex _buffer_map_lock; // buffer block map BufferMap _buffer_map; + std::unordered_map> _arrow_flight_query_buffers; // lock for timeout map std::mutex _timeout_lock; diff --git a/be/src/service/arrow_flight/arrow_flight_batch_reader.cpp b/be/src/service/arrow_flight/arrow_flight_batch_reader.cpp index 0a367692550..6fba86b7366 100644 --- a/be/src/service/arrow_flight/arrow_flight_batch_reader.cpp +++ b/be/src/service/arrow_flight/arrow_flight_batch_reader.cpp @@ -23,6 +23,8 @@ #include #include +#include + #include "core/block/block.h" #include "format/arrow/arrow_block_convertor.h" #include "format/arrow/arrow_row_batch.h" @@ -38,6 +40,39 @@ namespace doris::flight { +namespace { +// Poll cancellation on the RPC thread: no background thread may retain ServerCallContext. +template +class FlightReadCallback : public DummyBrpcCallback { +public: + void call() override { + std::scoped_lock lock(_mutex); + _done = true; + _cv.notify_all(); + } + + void wait(const std::function& is_cancelled) { + std::unique_lock lock(_mutex); + while (!_done) { + if (is_cancelled()) { + lock.unlock(); + brpc::StartCancel(this->call_id_); + this->join(); + return; + } + _cv.wait_for(lock, std::chrono::milliseconds(20)); + } + lock.unlock(); + this->join(); + } + +private: + std::mutex _mutex; + std::condition_variable _cv; + bool _done = false; +}; +} // namespace + ArrowFlightBatchReaderBase::ArrowFlightBatchReaderBase( const std::shared_ptr& statement) : _statement(statement) {} @@ -55,7 +90,46 @@ arrow::Status ArrowFlightBatchReaderBase::_return_invalid_status(const std::stri return arrow::Status::Invalid(status_msg); } +bool ArrowFlightBatchReaderBase::is_cancelled() const { + return _closed.load() || (_is_cancelled && _is_cancelled()); +} + +void ArrowFlightBatchReaderBase::close(const Status& reason) { + if (_closed.exchange(true) || _eof.load()) { + return; + } + auto* env = ExecEnv::GetInstance(); + if (_statement->result_addr.hostname == BackendOptions::get_localhost() && + _statement->result_addr.port == config::brpc_port) { + env->result_mgr()->cancel_arrow_flight_query(_statement->query_id, reason); + } +} + +arrow::Status ArrowFlightBatchReaderBase::Close() { + close(Status::Cancelled("Arrow Flight result stream closed before EOF")); + return arrow::Status::OK(); +} + +arrow::Status ArrowFlightBatchReaderBase::ReadNext(std::shared_ptr* out) { + *out = nullptr; + if (is_cancelled()) { + (void)Close(); + return arrow::Status::Cancelled("Arrow Flight fetch cancelled"); + } + auto status = [&]() -> arrow::Status { + RETURN_ARROW_STATUS_IF_CATCH_EXCEPTION(ReadNextImpl(out)); + }(); + if (!status.ok()) { + close(to_doris_status(status)); + } else if (!*out) { + _eof = true; + } + return status; +} + ArrowFlightBatchReaderBase::~ArrowFlightBatchReaderBase() { + // Transport errors can destroy a stream without another ReadNext/Close call. + (void)ArrowFlightBatchReaderBase::Close(); LOG(INFO) << fmt::format( "ArrowFlightBatchReader finished, packet_seq={}, result_addr={}:{}, finistId={}, " "convert_arrow_batch_timer={}, deserialize_block_timer={}, peak_memory_usage={}", @@ -74,7 +148,7 @@ ArrowFlightBatchLocalReader::ArrowFlightBatchLocalReader( } arrow::Result> ArrowFlightBatchLocalReader::Create( - const std::shared_ptr& statement) { + const std::shared_ptr& statement, std::function is_cancelled) { DCHECK(statement->result_addr.hostname == BackendOptions::get_localhost()); std::shared_ptr arrow_buffer; RETURN_ARROW_STATUS_IF_ERROR( @@ -86,6 +160,7 @@ arrow::Result> ArrowFlightBatchLoca std::shared_ptr mem_tracker = arrow_buffer->mem_tracker(); std::shared_ptr result( new ArrowFlightBatchLocalReader(statement, schema, mem_tracker)); + result->_is_cancelled = std::move(is_cancelled); arrow_buffer->get_timezone(result->_timezone_obj); return result; } @@ -99,10 +174,16 @@ arrow::Status ArrowFlightBatchLocalReader::ReadNextImpl(std::shared_ptrresult_mgr()->find_buffer(tid, arrow_buffer)); std::shared_ptr result; - auto st = arrow_buffer->get_arrow_batch(&result); - st.prepend("ArrowFlightBatchLocalReader fetch arrow data failed"); - ARROW_RETURN_NOT_OK(to_arrow_status(st)); - if (result == nullptr) { + bool eos = false; + while (!result && !eos) { + if (is_cancelled()) { + return arrow::Status::Cancelled("Arrow Flight fetch cancelled"); + } + auto st = arrow_buffer->get_arrow_batch(&result, &eos); + st.prepend("ArrowFlightBatchLocalReader fetch arrow data failed"); + ARROW_RETURN_NOT_OK(to_arrow_status(st)); + } + if (eos) { // eof, normal path end return arrow::Status::OK(); } @@ -110,8 +191,8 @@ arrow::Status ArrowFlightBatchLocalReader::ReadNextImpl(std::shared_ptr* out) { - RETURN_ARROW_STATUS_IF_CATCH_EXCEPTION(ReadNextImpl(out)); -} - ArrowFlightBatchRemoteReader::ArrowFlightBatchRemoteReader( const std::shared_ptr& statement, const std::shared_ptr& stub) @@ -138,7 +215,7 @@ ArrowFlightBatchRemoteReader::ArrowFlightBatchRemoteReader( } arrow::Result> ArrowFlightBatchRemoteReader::Create( - const std::shared_ptr& statement) { + const std::shared_ptr& statement, std::function is_cancelled) { std::shared_ptr stub = ExecEnv::GetInstance()->brpc_internal_client_cache()->get_client( statement->result_addr); @@ -153,6 +230,7 @@ arrow::Result> ArrowFlightBatchRem std::shared_ptr result( new ArrowFlightBatchRemoteReader(statement, stub)); + result->_is_cancelled = std::move(is_cancelled); ARROW_RETURN_NOT_OK(result->init_schema()); return result; } @@ -163,17 +241,20 @@ arrow::Status ArrowFlightBatchRemoteReader::_fetch_schema() { auto* pfinst_id = request->mutable_finst_id(); pfinst_id->set_hi(_statement->query_id.hi); pfinst_id->set_lo(_statement->query_id.lo); - auto callback = DummyBrpcCallback::create_shared(); + auto callback = std::make_shared>(); auto closure = AutoReleaseClosure< PFetchArrowFlightSchemaRequest, - DummyBrpcCallback>::create_unique(request, callback); + FlightReadCallback>::create_unique(request, callback); callback->cntl_->set_timeout_ms(config::arrow_flight_reader_brpc_controller_timeout_ms); callback->cntl_->ignore_eovercrowded(); _brpc_stub->fetch_arrow_flight_schema(closure->cntl_.get(), closure->request_.get(), closure->response_.get(), closure.get()); closure.release(); - callback->join(); + callback->wait([this] { return is_cancelled(); }); + if (is_cancelled()) { + return arrow::Status::Cancelled("Arrow Flight fetch cancelled"); + } if (callback->cntl_->Failed()) { if (!ExecEnv::GetInstance()->brpc_internal_client_cache()->available( @@ -213,17 +294,20 @@ arrow::Status ArrowFlightBatchRemoteReader::_fetch_data() { auto* pfinst_id = request->mutable_finst_id(); pfinst_id->set_hi(_statement->query_id.hi); pfinst_id->set_lo(_statement->query_id.lo); - auto callback = DummyBrpcCallback::create_shared(); + auto callback = std::make_shared>(); auto closure = AutoReleaseClosure< PFetchArrowDataRequest, - DummyBrpcCallback>::create_unique(request, callback); + FlightReadCallback>::create_unique(request, callback); callback->cntl_->set_timeout_ms(config::arrow_flight_reader_brpc_controller_timeout_ms); callback->cntl_->ignore_eovercrowded(); _brpc_stub->fetch_arrow_data(closure->cntl_.get(), closure->request_.get(), closure->response_.get(), closure.get()); closure.release(); - callback->join(); + callback->wait([this] { return is_cancelled(); }); + if (is_cancelled()) { + return arrow::Status::Cancelled("Arrow Flight fetch cancelled"); + } if (callback->cntl_->Failed()) { if (!ExecEnv::GetInstance()->brpc_internal_client_cache()->available( @@ -318,8 +402,4 @@ arrow::Status ArrowFlightBatchRemoteReader::ReadNextImpl(std::shared_ptr* out) { - RETURN_ARROW_STATUS_IF_CATCH_EXCEPTION(ReadNextImpl(out)); -} - } // namespace doris::flight diff --git a/be/src/service/arrow_flight/arrow_flight_batch_reader.h b/be/src/service/arrow_flight/arrow_flight_batch_reader.h index f72a35a8479..d8dc6618ebe 100644 --- a/be/src/service/arrow_flight/arrow_flight_batch_reader.h +++ b/be/src/service/arrow_flight/arrow_flight_batch_reader.h @@ -20,6 +20,7 @@ #include #include +#include #include #include @@ -48,10 +49,19 @@ class ArrowFlightBatchReaderBase : public arrow::RecordBatchReader { public: // RecordBatchReader force override [[nodiscard]] std::shared_ptr schema() const override; + arrow::Status Close() override; + arrow::Status ReadNext(std::shared_ptr* out) override; protected: ArrowFlightBatchReaderBase(const std::shared_ptr& statement); ~ArrowFlightBatchReaderBase() override; + virtual arrow::Status ReadNextImpl(std::shared_ptr* out) = 0; + bool is_cancelled() const; + void close(const Status& reason); + std::function _is_cancelled; + std::atomic _closed {false}; + std::atomic _eof {false}; + arrow::Status _return_invalid_status(const std::string& msg); std::shared_ptr _statement; @@ -67,34 +77,33 @@ protected: class ArrowFlightBatchLocalReader : public ArrowFlightBatchReaderBase { public: static arrow::Result> Create( - const std::shared_ptr& statement); - - arrow::Status ReadNext(std::shared_ptr* out) override; + const std::shared_ptr& statement, + std::function is_cancelled = {}); private: ArrowFlightBatchLocalReader(const std::shared_ptr& statement, const std::shared_ptr& schema, const std::shared_ptr& mem_tracker); - arrow::Status ReadNextImpl(std::shared_ptr* out); + arrow::Status ReadNextImpl(std::shared_ptr* out) override; }; class ArrowFlightBatchRemoteReader : public ArrowFlightBatchReaderBase { public: static arrow::Result> Create( - const std::shared_ptr& statement); + const std::shared_ptr& statement, + std::function is_cancelled = {}); // create arrow RecordBatchReader must initialize the schema. // so when creating arrow RecordBatchReader, fetch result data once, // which will return Block and some necessary information, and extract arrow schema from Block. arrow::Status init_schema(); - arrow::Status ReadNext(std::shared_ptr* out) override; private: ArrowFlightBatchRemoteReader(const std::shared_ptr& statement, const std::shared_ptr& stub); - arrow::Status ReadNextImpl(std::shared_ptr* out); + arrow::Status ReadNextImpl(std::shared_ptr* out) override; arrow::Status _fetch_schema(); arrow::Status _fetch_data(); diff --git a/be/src/service/arrow_flight/flight_sql_service.cpp b/be/src/service/arrow_flight/flight_sql_service.cpp index 26129559b2f..34acc0d6008 100644 --- a/be/src/service/arrow_flight/flight_sql_service.cpp +++ b/be/src/service/arrow_flight/flight_sql_service.cpp @@ -87,11 +87,15 @@ public: if (statement->result_addr.hostname == BackendOptions::get_localhost() && statement->result_addr.port == config::brpc_port) { std::shared_ptr reader; - ARROW_ASSIGN_OR_RAISE(reader, ArrowFlightBatchLocalReader::Create(statement)); + ARROW_ASSIGN_OR_RAISE( + reader, ArrowFlightBatchLocalReader::Create( + statement, [&context] { return context.is_cancelled(); })); return std::make_unique(reader); } else { std::shared_ptr reader; - ARROW_ASSIGN_OR_RAISE(reader, ArrowFlightBatchRemoteReader::Create(statement)); + ARROW_ASSIGN_OR_RAISE( + reader, ArrowFlightBatchRemoteReader::Create( + statement, [&context] { return context.is_cancelled(); })); return std::make_unique(reader); } } diff --git a/be/src/service/internal_service.cpp b/be/src/service/internal_service.cpp index b4b4e76537c..9d1c864eb1a 100644 --- a/be/src/service/internal_service.cpp +++ b/be/src/service/internal_service.cpp @@ -692,6 +692,11 @@ void PInternalService::cancel_plan_fragment(google::protobuf::RpcController* /*c LOG(INFO) << fmt::format("Cancel query {}, reason: {}", print_id(query_id), actual_cancel_status.to_string()); _exec_env->fragment_mgr()->cancel_query(query_id, actual_cancel_status); + // LIMIT_REACH/FINISHED can leave valid unread results; only an abort may discard them. + if (!actual_cancel_status.ok() && !actual_cancel_status.is() && + !actual_cancel_status.is()) { + _exec_env->result_mgr()->cancel_arrow_flight_buffers(query_id, actual_cancel_status); + } // TODO: the logic seems useless, cancel only return Status::OK. remove it st.to_protobuf(result->mutable_status()); @@ -731,6 +736,7 @@ void PInternalService::fetch_arrow_data(google::protobuf::RpcController* control auto st = ExecEnv::GetInstance()->result_mgr()->find_buffer(unique_id, arrow_buffer); if (!st.ok()) { LOG(WARNING) << "Result buffer not found! Query ID: " << print_id(unique_id); + ctx->on_failure(st); return; } if (st = arrow_buffer->get_batch(ctx); !st.ok()) { diff --git a/be/test/service/arrow_flight/arrow_flight_cancel_test.cpp b/be/test/service/arrow_flight/arrow_flight_cancel_test.cpp new file mode 100644 index 00000000000..b1830c17b34 --- /dev/null +++ b/be/test/service/arrow_flight/arrow_flight_cancel_test.cpp @@ -0,0 +1,415 @@ +// 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 +#include +#include +#include + +#include +#include + +#include "exec/pipeline/dependency.h" +#include "load/channel/load_stream_mgr.h" +#include "runtime/result_buffer_mgr.h" +#include "service/arrow_flight/arrow_flight_batch_reader.h" +#include "service/arrow_flight/flight_sql_service.h" +#include "service/backend_options.h" +#include "service/internal_service.h" +#include "testutil/column_helper.h" +#include "testutil/mock/mock_runtime_state.h" +#include "util/brpc_client_cache.h" +#include "util/debug/leakcheck_disabler.h" +#include "util/dns_cache.h" + +namespace doris::flight { + +class ArrowFlightCancelTest : public testing::Test { +protected: + void SetUp() override { + _previous_mgr = std::exchange(ExecEnv::GetInstance()->_result_mgr, &_mgr); + _state._batch_size = 1; + _id.hi = 1; + _id.lo = 2; + auto schema = arrow::schema({arrow::field("value", arrow::int64())}); + std::shared_ptr buffer; + ASSERT_TRUE(_mgr.create_sender(_id, 16, &buffer, &_state, true, schema).ok()); + _buffer = std::dynamic_pointer_cast(buffer); + _dep = Dependency::create_shared(0, 0, "Result", true); + _buffer->set_dependency(_state.fragment_instance_id(), _dep); + TNetworkAddress address; + address.hostname = BackendOptions::get_localhost(); + address.port = config::brpc_port; + _statement = std::make_shared(_id, address, "select value"); + } + void TearDown() override { + ExecEnv::GetInstance()->_result_mgr = _previous_mgr; + // The constructor registers a process-wide result_buffer_block_count hook that captures + // `this`; deregister it before the fixture destroys the manager. + _mgr.stop(); + } + void fill_buffer() { + auto block = std::make_shared(ColumnHelper::create_block({1, 2})); + ASSERT_TRUE(_buffer->add_batch(&_state, block).ok()); + ASSERT_FALSE(_dep->ready()); + } + bool registered() { + std::shared_ptr found; + return _mgr.find_buffer(_id, found).ok(); + } + MockRuntimeState _state; + ResultBufferMgr _mgr; + ResultBufferMgr* _previous_mgr = nullptr; + TUniqueId _id; + std::shared_ptr _buffer; + std::shared_ptr _dep; + std::shared_ptr _statement; +}; + +TEST_F(ArrowFlightCancelTest, EarlyCloseReleasesBackpressuredBuffer) { + fill_buffer(); + auto result = ArrowFlightBatchLocalReader::Create(_statement); + ASSERT_TRUE(result.ok()) << result.status(); + auto reader = *result; + ASSERT_TRUE(reader->Close().ok()); + EXPECT_FALSE(registered()); + EXPECT_TRUE(_dep->ready()); + EXPECT_TRUE(_state.get_query_ctx()->is_cancelled()); + std::shared_ptr block; + bool eos = false; + EXPECT_FALSE(_buffer->get_arrow_batch(&block, &eos).ok()); + EXPECT_TRUE(reader->Close().ok()); +} + +TEST_F(ArrowFlightCancelTest, DestructionReleasesAbandonedBuffer) { + fill_buffer(); + { + auto reader = ArrowFlightBatchLocalReader::Create(_statement); + ASSERT_TRUE(reader.ok()) << reader.status(); + } + EXPECT_FALSE(registered()); + EXPECT_TRUE(_dep->ready()); + EXPECT_TRUE(_state.get_query_ctx()->is_cancelled()); +} + +TEST_F(ArrowFlightCancelTest, CloseReleasesSiblingBuffersAfterContextExpires) { + TUniqueId sibling_id = _id; + ++sibling_id.lo; + std::shared_ptr sibling; + ASSERT_TRUE(_mgr.create_sender(sibling_id, 16, &sibling, &_state, true, + arrow::schema({arrow::field("value", arrow::int64())})) + .ok()); + auto arrow_sibling = std::dynamic_pointer_cast(sibling); + auto block = std::make_shared(ColumnHelper::create_block({1, 2})); + ASSERT_TRUE(arrow_sibling->add_batch(&_state, block).ok()); + auto reader = ArrowFlightBatchLocalReader::Create(_statement); + ASSERT_TRUE(reader.ok()) << reader.status(); + _state._query_ctx_uptr.reset(); + _state._query_ctx = nullptr; + ASSERT_TRUE((*reader)->Close().ok()); + EXPECT_FALSE(registered()); + std::shared_ptr found; + EXPECT_FALSE(_mgr.find_buffer(sibling_id, found).ok()); + EXPECT_TRUE(arrow_sibling->_result_batch_queue.empty()); +} + +TEST_F(ArrowFlightCancelTest, NormalEofIsNotCancellation) { + bool fully_closed = false; + ASSERT_TRUE(_buffer->close(_state.fragment_instance_id(), Status::OK(), 0, fully_closed).ok()); + auto result = ArrowFlightBatchLocalReader::Create(_statement); + ASSERT_TRUE(result.ok()) << result.status(); + auto reader = *result; + std::shared_ptr batch; + ASSERT_TRUE(reader->ReadNext(&batch).ok()); + EXPECT_EQ(batch, nullptr); + EXPECT_TRUE(reader->Close().ok()); + std::shared_ptr block; + bool eos = false; + EXPECT_TRUE(_buffer->get_arrow_batch(&block, &eos).ok()); + EXPECT_TRUE(eos); + EXPECT_FALSE(_state.get_query_ctx()->is_cancelled()); +} + +TEST_F(ArrowFlightCancelTest, CancellationInterruptsEmptyBufferWait) { + std::atomic cancelled = false; + std::promise checked; + std::atomic checks = 0; + auto result = ArrowFlightBatchLocalReader::Create(_statement, [&] { + if (checks.fetch_add(1) == 3) { + checked.set_value(); + } + return cancelled.load(); + }); + ASSERT_TRUE(result.ok()) << result.status(); + auto reader = *result; + auto fetch = std::async(std::launch::async, [&] { + std::shared_ptr batch; + return reader->ReadNext(&batch); + }); + EXPECT_EQ(checked.get_future().wait_for(std::chrono::seconds(2)), std::future_status::ready); + cancelled = true; + const auto ready = fetch.wait_for(std::chrono::seconds(2)); + // Ensure a broken implementation fails instead of hanging the test runner. + if (ready != std::future_status::ready) { + _buffer->cancel(Status::Cancelled("test cleanup")); + } + EXPECT_EQ(ready, std::future_status::ready); + EXPECT_FALSE(fetch.get().ok()); + EXPECT_FALSE(registered()); + EXPECT_TRUE(_state.get_query_ctx()->is_cancelled()); +} + +TEST_F(ArrowFlightCancelTest, EmptyBufferReadReturnsForRetry) { + auto fetch = std::async(std::launch::async, [&] { + auto block = std::make_shared(); + bool eos = true; + auto status = _buffer->get_arrow_batch(&block, &eos); + EXPECT_EQ(block, nullptr); + EXPECT_FALSE(eos); + return status; + }); + const auto ready = fetch.wait_for(std::chrono::seconds(2)); + // Bound the test even if an empty buffer incorrectly waits until query completion. + if (ready != std::future_status::ready) { + _buffer->cancel(Status::Cancelled("test cleanup")); + } + EXPECT_EQ(ready, std::future_status::ready); + EXPECT_TRUE(fetch.get().ok()); +} + +// GTest assertion macros inflate the metric of this wait, retry and EOF sequence. +TEST_F(ArrowFlightCancelTest, // NOLINT(readability-function-cognitive-complexity) + EmptyWaitsDoNotFinishReader) { + std::promise retried; + std::atomic checks = 0; + auto result = ArrowFlightBatchLocalReader::Create(_statement, [&] { + if (checks.fetch_add(1) == 3) { + retried.set_value(); + } + return false; + }); + ASSERT_TRUE(result.ok()) << result.status(); + auto reader = *result; + std::shared_ptr batch; + auto fetch = std::async(std::launch::async, [&] { return reader->ReadNext(&batch); }); + EXPECT_EQ(retried.get_future().wait_for(std::chrono::seconds(2)), std::future_status::ready); + EXPECT_EQ(fetch.wait_for(std::chrono::seconds(0)), std::future_status::timeout); + auto block = std::make_shared(ColumnHelper::create_block({1, 2})); + EXPECT_TRUE(_buffer->add_batch(&_state, block).ok()); + bool fully_closed = false; + EXPECT_TRUE(_buffer->close(_state.fragment_instance_id(), Status::OK(), 2, fully_closed).ok()); + const auto ready = fetch.wait_for(std::chrono::seconds(2)); + if (ready != std::future_status::ready) { + _buffer->cancel(Status::Cancelled("test cleanup")); + } + EXPECT_EQ(ready, std::future_status::ready); + ASSERT_TRUE(fetch.get().ok()); + ASSERT_NE(batch, nullptr); + EXPECT_EQ(batch->num_rows(), 2); + ASSERT_TRUE(reader->ReadNext(&batch).ok()); + EXPECT_EQ(batch, nullptr); + EXPECT_TRUE(reader->Close().ok()); + EXPECT_FALSE(_state.get_query_ctx()->is_cancelled()); +} + +TEST_F(ArrowFlightCancelTest, ConversionFailureCancelsQuery) { + auto block = std::make_shared(ColumnHelper::create_block({"bad"})); + ASSERT_TRUE(_buffer->add_batch(&_state, block).ok()); + auto result = ArrowFlightBatchLocalReader::Create(_statement); + ASSERT_TRUE(result.ok()) << result.status(); + std::shared_ptr batch; + EXPECT_FALSE((*result)->ReadNext(&batch).ok()); + EXPECT_FALSE(registered()); + EXPECT_TRUE(_state.get_query_ctx()->is_cancelled()); +} + +TEST_F(ArrowFlightCancelTest, RealFlightRpcCancellationCleansUpQuery) { + auto created = FlightSqlServer::create(); + ASSERT_TRUE(created.ok()) << created.status(); + auto server = *created; + ASSERT_TRUE(server->init(0).ok()); + auto location = arrow::flight::Location::ForGrpcTcp("127.0.0.1", server->port()); + ASSERT_TRUE(location.ok()) << location.status(); + auto client_options = arrow::flight::FlightClientOptions::Defaults(); + client_options.generic_options.emplace_back("grpc.enable_http_proxy", 0); + auto connected = arrow::flight::FlightClient::Connect(*location, client_options); + ASSERT_TRUE(connected.ok()) << connected.status(); + auto client = std::move(*connected); + auto handle = arrow::flight::sql::CreateStatementQueryTicket( + print_id(_id) + "&" + BackendOptions::get_localhost() + "&" + + std::to_string(config::brpc_port) + "&"); + ASSERT_TRUE(handle.ok()) << handle.status(); + arrow::flight::FlightCallOptions options; + options.timeout = std::chrono::seconds(5); + auto fetched = client->DoGet(options, arrow::flight::Ticket {*handle}); + ASSERT_TRUE(fetched.ok()) << fetched.status(); + auto reader = std::move(*fetched); + reader->Cancel(); + for (int i = 0; i < 200 && (registered() || !_state.get_query_ctx()->is_cancelled()); ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + EXPECT_FALSE(registered()); + EXPECT_TRUE(_state.get_query_ctx()->is_cancelled()); + _buffer->cancel(Status::Cancelled("test cleanup")); + reader.reset(); + EXPECT_TRUE(client->Close().ok()); + EXPECT_TRUE(server->join().ok()); +} + +class ArrowFlightRemoteCancelTest : public ArrowFlightCancelTest { +protected: + void SetUp() override { + ArrowFlightCancelTest::SetUp(); + static DNSCache dns_cache; + _previous_dns_cache = std::exchange(ExecEnv::GetInstance()->_dns_cache, &dns_cache); + _previous_cache = std::exchange(ExecEnv::GetInstance()->_internal_client_cache, &_cache); + auto heavy = std::exchange(config::brpc_heavy_work_pool_threads, 1); + auto light = std::exchange(config::brpc_light_work_pool_threads, 1); + auto flight = std::exchange(config::brpc_arrow_flight_work_pool_threads, 2); + _previous_load_mgr = std::move(ExecEnv::GetInstance()->_load_stream_mgr); + ExecEnv::GetInstance()->_load_stream_mgr = std::make_unique(1); + _service = std::make_unique(ExecEnv::GetInstance()); + config::brpc_heavy_work_pool_threads = heavy; + config::brpc_light_work_pool_threads = light; + config::brpc_arrow_flight_work_pool_threads = flight; + ASSERT_EQ(_server.AddService(_service.get(), brpc::SERVER_DOESNT_OWN_SERVICE), 0); + { + // brpc registers process-lifetime built-in services on first startup. + debug::ScopedLeakCheckDisabler disable_lsan; + ASSERT_EQ(_server.Start("127.0.0.1:0", nullptr), 0); + } + _statement->result_addr.hostname = "127.0.0.1"; + _statement->result_addr.port = _server.listen_address().port; + _fragment_mgr = + std::make_unique(_cancel_status, ExecEnv::GetInstance()); + _fragment_mgr->stop(); + _previous_fragment_mgr = + std::exchange(ExecEnv::GetInstance()->_fragment_mgr, _fragment_mgr.get()); + } + void TearDown() override { + ExecEnv::GetInstance()->_dns_cache = _previous_dns_cache; + ExecEnv::GetInstance()->_fragment_mgr = _previous_fragment_mgr; + _buffer->cancel(Status::Cancelled("test cleanup")); + _server.Stop(0); + _server.Join(); + _service.reset(); + ExecEnv::GetInstance()->_load_stream_mgr = std::move(_previous_load_mgr); + ExecEnv::GetInstance()->_internal_client_cache = _previous_cache; + ArrowFlightCancelTest::TearDown(); + } + BrpcClientCache _cache; + BrpcClientCache* _previous_cache = nullptr; + std::unique_ptr _previous_load_mgr; + std::unique_ptr _service; + brpc::Server _server; + DNSCache* _previous_dns_cache = nullptr; + Status _cancel_status; + std::unique_ptr _fragment_mgr; + FragmentMgr* _previous_fragment_mgr = nullptr; +}; + +TEST_F(ArrowFlightRemoteCancelTest, CancellationInterruptsPendingBrpcFetch) { + std::atomic cancelled = false; + auto result = + ArrowFlightBatchRemoteReader::Create(_statement, [&] { return cancelled.load(); }); + ASSERT_TRUE(result.ok()) << result.status(); + auto fetch = std::async(std::launch::async, [&] { + std::shared_ptr batch; + return (*result)->ReadNext(&batch); + }); + // Wait until the server has queued the fetch; cancelling before that misses the blocking path. + bool waiting = false; + for (int i = 0; i < 200; ++i) { + { + std::scoped_lock lock(_buffer->_lock); + waiting = !_buffer->_waiting_rpc.empty(); + } + if (waiting) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + EXPECT_TRUE(waiting); + cancelled = true; + auto ready = fetch.wait_for(std::chrono::seconds(3)); + if (ready != std::future_status::ready) { + _buffer->cancel(Status::Cancelled("test cleanup")); + } + EXPECT_EQ(ready, std::future_status::ready); + // Releasing the remote result buffer needs an FE cancellation route; only the interrupted + // BRPC wait is checked here. + EXPECT_FALSE(fetch.get().ok()); +} + +TEST_F(ArrowFlightRemoteCancelTest, FetchAfterCancellationStillCompletesRpc) { + auto reader = ArrowFlightBatchRemoteReader::Create(_statement); + ASSERT_TRUE(reader.ok()) << reader.status(); + auto stub = _cache.get_client(_statement->result_addr); + ASSERT_NE(stub, nullptr); + { + // An aborting fragment cancellation must also drop the query's unread Flight buffers. + PCancelPlanFragmentRequest request; + request.mutable_finst_id()->set_hi(_id.hi); + request.mutable_finst_id()->set_lo(_id.lo); + request.mutable_query_id()->set_hi(_state.query_id().hi); + request.mutable_query_id()->set_lo(_state.query_id().lo); + Status::Cancelled("Flight stream aborted").to_protobuf(request.mutable_cancel_status()); + PCancelPlanFragmentResult response; + brpc::Controller controller; + controller.set_timeout_ms(1000); + stub->cancel_plan_fragment(&controller, &request, &response, nullptr); + ASSERT_FALSE(controller.Failed()) << controller.ErrorText(); + EXPECT_TRUE(Status::create(response.status()).ok()); + } + EXPECT_FALSE(_cancel_status.ok()); + EXPECT_FALSE(registered()); + EXPECT_TRUE(_state.get_query_ctx()->is_cancelled()); + PFetchArrowDataRequest request; + request.mutable_finst_id()->set_hi(_id.hi); + request.mutable_finst_id()->set_lo(_id.lo); + PFetchArrowDataResult response; + brpc::Controller controller; + controller.set_timeout_ms(1000); + stub->fetch_arrow_data(&controller, &request, &response, nullptr); + EXPECT_FALSE(controller.Failed()) << controller.ErrorText(); + EXPECT_TRUE(response.has_status()); + EXPECT_FALSE(Status::create(response.status()).ok()); +} + +TEST_F(ArrowFlightRemoteCancelTest, SuccessfulFragmentStopPreservesUnreadResults) { + fill_buffer(); + auto stub = _cache.get_client(_statement->result_addr); + for (const auto& reason : {Status::Error("limit reached"), + Status::Error("finished")}) { + PCancelPlanFragmentRequest request; + request.mutable_finst_id()->set_hi(_id.hi); + request.mutable_finst_id()->set_lo(_id.lo); + request.mutable_query_id()->set_hi(_state.query_id().hi); + request.mutable_query_id()->set_lo(_state.query_id().lo); + reason.to_protobuf(request.mutable_cancel_status()); + PCancelPlanFragmentResult response; + brpc::Controller controller; + controller.set_timeout_ms(1000); + stub->cancel_plan_fragment(&controller, &request, &response, nullptr); + ASSERT_FALSE(controller.Failed()) << controller.ErrorText(); + EXPECT_TRUE(Status::create(response.status()).ok()); + EXPECT_TRUE(registered()); + EXPECT_FALSE(_buffer->_result_batch_queue.empty()); + } +} + +} // namespace doris::flight