[fix](be) Cancel Arrow Flight queries on the result BE when a result stream closes early (#69945)

### What problem does this PR solve?

Issue Number: None

Related PR: #68630, #67644

Problem Summary:

#### In short

This ports the BE half of branch-4.x's #68630 to master. When an Arrow
Flight SQL client cancels or closes a result stream before its end,
master keeps the query running: the BE thread serving the stream waits
on the result buffer without ever checking whether the client is still
there, and the unread result blocks stay in memory, holding back the
query, until the query times out or the delayed buffer cleanup runs
minutes later. With this PR the BE notices the cancelled stream within
20 ms, and when the stream is served by the BE that holds its result
(the usual case), that BE cancels the query and frees all of the query's
Flight result buffers there. The FE route that also reaches the query's
other BEs (`cancelFlightQuery`, which changes `FrontendService.thrift`)
comes in a follow-up PR.

#### Background

- A Flight SQL query leaves its result on the BEs. `GetFlightInfo`
returns one endpoint per result buffer, and the client reads each
endpoint with a `DoGet` stream. Since #67504 the FE closes the
coordinator as soon as it has returned the `FlightInfo`, so it no longer
cancels anything on behalf of an open stream.
- `DoGet` is served by `ArrowFlightBatchLocalReader` when the result
buffer is on the same BE, or by `ArrowFlightBatchRemoteReader`, which
pulls the blocks from the result BE with the `fetch_arrow_data` BRPC.
- The result buffer (`ArrowFlightResultBlockBuffer`) is keyed by the
query id when parallel result sink is on (the default) and by the
fragment instance id when it is off. When it is full, the result sink
stops and the query waits for the reader.
- `ResultBufferMgr` drops a buffer that nobody removed at
`execution_timeout + 5 s`, or `result_buffer_cancelled_interval_time`
(300 s by default) after the fragments writing it have finished,
whichever comes first. Finishing with an error does not clear the unread
blocks. `cancel_plan_fragment` is the BRPC the FE uses to cancel a query
on a BE.

#### The problem, and what it cost

- **The local reader never looks at the client.** `get_arrow_batch`
loops on 20 ms waits until a block arrives, the buffer is closed or it
fails. It has no access to `ServerCallContext::is_cancelled()`, so a
`DoGet` cancelled while the buffer is empty keeps a Flight server thread
waiting until the query produces a block or ends.
- **Closing a stream cancels nothing.** The readers do not override
`RecordBatchReader::Close()`, and their destructor only logs. A client
that stops reading (a `LIMIT` reached on the client side, a cancelled
cursor, a dropped connection) leaves the query running and its unread
blocks in the buffer. When the result does not fit in the buffer, the
full buffer back-pressures the result sink and the fragments, with their
memory, wait until the query times out; the unread blocks go only when
the buffer is dropped.
- **The remote reader cannot be interrupted.** It waits for
`fetch_arrow_data` with `brpc::Join`, which returns only when the result
BE answers or `arrow_flight_reader_brpc_controller_timeout_ms` (300 s by
default) passes.
- **An FE cancel leaves the results behind.** `cancel_plan_fragment`
cancels the fragments but not the query's Flight buffers, so their
unread blocks stay until the delayed cleanup 300 s later.
- **A fetch for a removed buffer never completes.** When
`fetch_arrow_data` cannot find the buffer, it logs and returns without
running `done`, so the caller waits for its BRPC timeout instead of
getting an error. #67644 fixes the same line.

#### How this PR fixes it

- **Bounded waits, cancellation checks in the reader.**
`get_arrow_batch(result, eos)` waits at most 20 ms and then returns OK
with no block and `eos = false`; `eos` now marks the real end. The local
reader retries in a loop and checks `is_cancelled()` before each wait.
`DoGetStatement` passes `context.is_cancelled()` to both readers.
- **The reader closes itself.** `ArrowFlightBatchReaderBase` now
implements `ReadNext` (which calls each reader's `ReadNextImpl`) and
`Close()`. A stream counts as aborted when it is cancelled, closed
before EOF, fails while reading, or is destroyed before EOF; a stream
that reached EOF is never treated as aborted.
- **Aborting a local stream cancels the query on this BE.**
`ResultBufferMgr` keeps an index from query id to the query's Flight
buffers, maintained under the buffer-map lock together with the map.
`cancel_arrow_flight_buffers` removes every Flight buffer of the query
from the map, then, outside the lock, fails each buffer (which wakes its
waiting fetches and releases the back-pressure) and cancels the
`QueryContext` it holds a weak reference to. Sibling buffers of the same
query are cleared even after their `QueryContext` has expired.
- **Remote waits can be interrupted.** The remote reader waits for its
BRPC with a callback that checks `is_cancelled()` every 20 ms and calls
`brpc::StartCancel` when the stream is gone.
- **`cancel_plan_fragment` drops the Flight buffers of an aborted
query.** It calls `cancel_arrow_flight_buffers` unless the reason is
`LIMIT_REACH` or `FINISHED`, which can leave valid unread results that
the client still has to read.
- **`fetch_arrow_data` answers when the buffer is gone.** It completes
the RPC with the lookup error. The line is identical to the one #67644
adds, so the two PRs merge in either order.

Differences from branch-4.x: the two consistency checks between the
index and the buffer map are `DORIS_CHECK` (an `if` and a `DCHECK` in
4.x), a null check on the reader's memory tracker that both constructors
make unnecessary is not carried over, and the
`ArrowFlightResultBlockBuffer` constructor moved to the .cpp so that
`varrow_flight_result_writer.h`, which 13 TUs include, does not include
`runtime/query_context.h`; only the .cpp needs its definition. The new
lines also take a few clang-tidy fixes (the destructor calls
`ArrowFlightBatchReaderBase::Close()` explicitly instead of a virtual
call, `std::scoped_lock`), and the test starts its BRPC server inside
`ScopedLeakCheckDisabler`, like the other BE tests that start one,
because brpc registers process-lifetime built-in services on its first
start.

Not in this PR: the FE route. A remote reader's abort, and an abort on
one result BE, still do not reach the query's other BEs; their fragments
and buffers wait for the timeout as before. The follow-up adds
`FrontendService.cancelFlightQuery` (same signature as branch-4.x), the
BE broadcast to the FEs and the FE routing table, together with the FE
unit test and the `test_flight_cancel_cleanup` regression test.

| Situation | Before | After |
|---|---|---|
| Client cancels `DoGet` while the result buffer is empty (local reader)
| Server thread waits until the query produces data or ends | Returns
within 20 ms; the query is cancelled on this BE |
| Client closes a local stream before EOF | Query keeps running,
back-pressured until it times out if the result does not fit; unread
blocks kept for minutes | Query cancelled on this BE; all of its Flight
buffers there freed at once |
| Client cancels while a remote fetch is pending | Waits for the result
BE, up to 300 s | BRPC cancelled within 20 ms |
| FE cancels the query (`cancel_plan_fragment`, abort reason) | Unread
blocks kept until the delayed cleanup, 300 s later | Flight buffers
freed at once; `LIMIT_REACH` / `FINISHED` still keep them |
| `fetch_arrow_data` for a removed buffer | RPC never completes; caller
waits for its timeout | Completes at once with an error |
| Stream read to EOF, then closed | Nothing cancelled | Nothing
cancelled |
| New `DoGet` with the ticket of a stream aborted on its result BE |
Reads the remaining blocks, or EOS | Fails: the buffer is gone |

#### Results

- `ArrowFlightCancelTest` (9 tests, new): early close and destruction
free a back-pressured buffer and cancel the query; sibling buffers are
freed after their `QueryContext` expired; EOF followed by close is not a
cancellation; cancellation interrupts an empty-buffer wait; empty waits
return for retry and do not end the stream; a conversion error cancels
the query; a real Flight server sees a client-side `Cancel()` through
`ServerCallContext` and cleans up.
- `ArrowFlightRemoteCancelTest` (3 tests, new, against an in-process
BRPC `PInternalService`): cancellation interrupts a pending
`fetch_arrow_data`; an aborting `cancel_plan_fragment` frees the buffer
and a following fetch completes with an error inside its 1 s timeout;
`LIMIT_REACH` / `FINISHED` keep the unread results.
- On a local single-BE cluster, branch-4.x's regression
`test_flight_cancel_cleanup` (it aborts one endpoint after the first
batch, by closing it and by an explicit cancel, and requires the query's
pipelines to disappear within 20 s) fails on master with `Cancelled
Flight query retained pipelines: [...: 1 pipeline fragment contexts in
query ...]` and passes with this PR. The suite comes with the follow-up,
because on a multi-BE cluster it also needs the FE route. The 20
`arrow_flight_sql_p0` suites and the 14 `external_table_p0/remote_doris`
suites, which read results over Arrow Flight, pass before and after.
- The 12 new tests plus `ResultBufferMgrTest` (6) and
`ArrowResultBlockBufferTest` (4) pass 10 times in a row in one ASAN
process. Ten single-point mutations of the change (dropping the buffer
cleanup in `cancel_plan_fragment`, the `on_failure` call, the local
cancel in `close()`, the cancellation check in the local read loop or in
the BRPC wait, the sibling loop, the EOF flag, the `QueryContext`
cancel, the cancel in the destructor, the `eos = false` on an empty
wait) each turn their target tests red.

be/src/exec/sink/writer/varrow_flight_result_writer.cpp

@@ -21,8 +21,11 @@
 #include <gen_cpp/PaloInternalService_types.h>
 #include <gen_cpp/internal_service.pb.h>
 
+#include <utility>
+
 #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<arrow::Schema>*
                                              print_id(_fragment_id), _status));
 }
 
-Status ArrowFlightResultBlockBuffer::get_arrow_batch(std::shared_ptr<Block>* result) {
+ArrowFlightResultBlockBuffer::ArrowFlightResultBlockBuffer(TUniqueId id, RuntimeState* state,
+                                                           std::shared_ptr<arrow::Schema> schema,
+                                                           int buffer_size)
+        : ResultBlockBuffer<GetArrowResultBatchCtx>(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<QueryContext> {}) {
+    _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<Block>* result, bool* eos) {
+    *result = nullptr;
+    *eos = false;
     std::unique_lock<std::mutex> 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<Block>* res
 
     // normal path end
     if (_is_close) {
+        *eos = true;
         if (!_status.ok()) {
             return _status;
         }

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<GetArrowResu
 public:
     using ResultType = Block;
     ArrowFlightResultBlockBuffer(TUniqueId id, RuntimeState* state,
-                                 std::shared_ptr<arrow::Schema> schema, int buffer_size)
-            : ResultBlockBuffer<GetArrowResultBatchCtx>(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<arrow::Schema> schema, int buffer_size);
     ~ArrowFlightResultBlockBuffer() override = default;
-    Status get_arrow_batch(std::shared_ptr<Block>* result);
+    // A timed wait may return OK with no block and eos=false; the caller must retry.
+    Status get_arrow_batch(std::shared_ptr<Block>* 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>* 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<QueryContext> _query_ctx;
 };
 
 class VArrowFlightResultWriter final : public ResultWriter {

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<std::shared_mutex> 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<ArrowFlightResultBlockBuffer>(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<ArrowFlightResultBlockBuffer> 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<std::shared_ptr<ArrowFlightResultBlockBuffer>> buffers;
+    {
+        std::unique_lock<std::shared_mutex> 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<ArrowFlightResultBlockBuffer>(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<std::mutex> l(_timeout_lock);
     auto iter = _timeout_map.find(cancel_time);

be/src/runtime/result_buffer_mgr.h

@@ -27,6 +27,7 @@
 #include <mutex>
 #include <shared_mutex>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 #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<TUniqueId, std::unordered_set<TUniqueId>> _arrow_flight_query_buffers;
 
     // lock for timeout map
     std::mutex _timeout_lock;

be/src/service/arrow_flight/arrow_flight_batch_reader.cpp

@@ -23,6 +23,8 @@
 #include <arrow/type.h>
 #include <gen_cpp/internal_service.pb.h>
 
+#include <condition_variable>
+
 #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 <typename Response>
+class FlightReadCallback : public DummyBrpcCallback<Response> {
+public:
+    void call() override {
+        std::scoped_lock lock(_mutex);
+        _done = true;
+        _cv.notify_all();
+    }
+
+    void wait(const std::function<bool()>& 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<QueryStatement>& 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<arrow::RecordBatch>* 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<std::shared_ptr<ArrowFlightBatchLocalReader>> ArrowFlightBatchLocalReader::Create(
-        const std::shared_ptr<QueryStatement>& statement) {
+        const std::shared_ptr<QueryStatement>& statement, std::function<bool()> is_cancelled) {
     DCHECK(statement->result_addr.hostname == BackendOptions::get_localhost());
     std::shared_ptr<ArrowFlightResultBlockBuffer> arrow_buffer;
     RETURN_ARROW_STATUS_IF_ERROR(
@@ -86,6 +160,7 @@ arrow::Result<std::shared_ptr<ArrowFlightBatchLocalReader>> ArrowFlightBatchLoca
     std::shared_ptr<MemTrackerLimiter> mem_tracker = arrow_buffer->mem_tracker();
     std::shared_ptr<ArrowFlightBatchLocalReader> 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_ptr<arrow::R
     RETURN_ARROW_STATUS_IF_ERROR(
             ExecEnv::GetInstance()->result_mgr()->find_buffer(tid, arrow_buffer));
     std::shared_ptr<Block> 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<arrow::R
     {
         // convert one batch
         SCOPED_ATOMIC_TIMER(&_convert_arrow_batch_timer);
-        st = ArrowFlightArrowBlockConvertor(_schema, _timezone_obj)
-                     .convert_to_arrow(*result, arrow::default_memory_pool(), out);
+        auto st = ArrowFlightArrowBlockConvertor(_schema, _timezone_obj)
+                          .convert_to_arrow(*result, arrow::default_memory_pool(), out);
         st.prepend("ArrowFlightBatchLocalReader convert block to arrow batch failed");
         ARROW_RETURN_NOT_OK(to_arrow_status(st));
     }
@@ -124,10 +205,6 @@ arrow::Status ArrowFlightBatchLocalReader::ReadNextImpl(std::shared_ptr<arrow::R
     return arrow::Status::OK();
 }
 
-arrow::Status ArrowFlightBatchLocalReader::ReadNext(std::shared_ptr<arrow::RecordBatch>* out) {
-    RETURN_ARROW_STATUS_IF_CATCH_EXCEPTION(ReadNextImpl(out));
-}
-
 ArrowFlightBatchRemoteReader::ArrowFlightBatchRemoteReader(
         const std::shared_ptr<QueryStatement>& statement,
         const std::shared_ptr<PBackendService_Stub>& stub)
@@ -138,7 +215,7 @@ ArrowFlightBatchRemoteReader::ArrowFlightBatchRemoteReader(
 }
 
 arrow::Result<std::shared_ptr<ArrowFlightBatchRemoteReader>> ArrowFlightBatchRemoteReader::Create(
-        const std::shared_ptr<QueryStatement>& statement) {
+        const std::shared_ptr<QueryStatement>& statement, std::function<bool()> is_cancelled) {
     std::shared_ptr<PBackendService_Stub> stub =
             ExecEnv::GetInstance()->brpc_internal_client_cache()->get_client(
                     statement->result_addr);
@@ -153,6 +230,7 @@ arrow::Result<std::shared_ptr<ArrowFlightBatchRemoteReader>> ArrowFlightBatchRem
 
     std::shared_ptr<ArrowFlightBatchRemoteReader> 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<PFetchArrowFlightSchemaResult>::create_shared();
+    auto callback = std::make_shared<FlightReadCallback<PFetchArrowFlightSchemaResult>>();
     auto closure = AutoReleaseClosure<
             PFetchArrowFlightSchemaRequest,
-            DummyBrpcCallback<PFetchArrowFlightSchemaResult>>::create_unique(request, callback);
+            FlightReadCallback<PFetchArrowFlightSchemaResult>>::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<PFetchArrowDataResult>::create_shared();
+        auto callback = std::make_shared<FlightReadCallback<PFetchArrowDataResult>>();
         auto closure = AutoReleaseClosure<
                 PFetchArrowDataRequest,
-                DummyBrpcCallback<PFetchArrowDataResult>>::create_unique(request, callback);
+                FlightReadCallback<PFetchArrowDataResult>>::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<arrow::
     return arrow::Status::OK();
 }
 
-arrow::Status ArrowFlightBatchRemoteReader::ReadNext(std::shared_ptr<arrow::RecordBatch>* out) {
-    RETURN_ARROW_STATUS_IF_CATCH_EXCEPTION(ReadNextImpl(out));
-}
-
 } // namespace doris::flight

be/src/service/arrow_flight/arrow_flight_batch_reader.h

@@ -20,6 +20,7 @@
 #include <cctz/time_zone.h>
 #include <gen_cpp/Types_types.h>
 
+#include <functional>
 #include <memory>
 #include <utility>
 
@@ -48,10 +49,19 @@ class ArrowFlightBatchReaderBase : public arrow::RecordBatchReader {
 public:
     // RecordBatchReader force override
     [[nodiscard]] std::shared_ptr<arrow::Schema> schema() const override;
+    arrow::Status Close() override;
+    arrow::Status ReadNext(std::shared_ptr<arrow::RecordBatch>* out) override;
 
 protected:
     ArrowFlightBatchReaderBase(const std::shared_ptr<QueryStatement>& statement);
     ~ArrowFlightBatchReaderBase() override;
+    virtual arrow::Status ReadNextImpl(std::shared_ptr<arrow::RecordBatch>* out) = 0;
+    bool is_cancelled() const;
+    void close(const Status& reason);
+    std::function<bool()> _is_cancelled;
+    std::atomic<bool> _closed {false};
+    std::atomic<bool> _eof {false};
+
     arrow::Status _return_invalid_status(const std::string& msg);
 
     std::shared_ptr<QueryStatement> _statement;
@@ -67,34 +77,33 @@ protected:
 class ArrowFlightBatchLocalReader : public ArrowFlightBatchReaderBase {
 public:
     static arrow::Result<std::shared_ptr<ArrowFlightBatchLocalReader>> Create(
-            const std::shared_ptr<QueryStatement>& statement);
-
-    arrow::Status ReadNext(std::shared_ptr<arrow::RecordBatch>* out) override;
+            const std::shared_ptr<QueryStatement>& statement,
+            std::function<bool()> is_cancelled = {});
 
 private:
     ArrowFlightBatchLocalReader(const std::shared_ptr<QueryStatement>& statement,
                                 const std::shared_ptr<arrow::Schema>& schema,
                                 const std::shared_ptr<MemTrackerLimiter>& mem_tracker);
 
-    arrow::Status ReadNextImpl(std::shared_ptr<arrow::RecordBatch>* out);
+    arrow::Status ReadNextImpl(std::shared_ptr<arrow::RecordBatch>* out) override;
 };
 
 class ArrowFlightBatchRemoteReader : public ArrowFlightBatchReaderBase {
 public:
     static arrow::Result<std::shared_ptr<ArrowFlightBatchRemoteReader>> Create(
-            const std::shared_ptr<QueryStatement>& statement);
+            const std::shared_ptr<QueryStatement>& statement,
+            std::function<bool()> 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<arrow::RecordBatch>* out) override;
 
 private:
     ArrowFlightBatchRemoteReader(const std::shared_ptr<QueryStatement>& statement,
                                  const std::shared_ptr<PBackendService_Stub>& stub);
 
-    arrow::Status ReadNextImpl(std::shared_ptr<arrow::RecordBatch>* out);
+    arrow::Status ReadNextImpl(std::shared_ptr<arrow::RecordBatch>* out) override;
     arrow::Status _fetch_schema();
     arrow::Status _fetch_data();
 

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<ArrowFlightBatchLocalReader> 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<arrow::flight::RecordBatchStream>(reader);
         } else {
             std::shared_ptr<ArrowFlightBatchRemoteReader> 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<arrow::flight::RecordBatchStream>(reader);
         }
     }

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<ErrorCode::LIMIT_REACH>() &&
+            !actual_cancel_status.is<ErrorCode::FINISHED>()) {
+            _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()) {

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 <arrow/flight/client.h>
+#include <arrow/flight/sql/server.h>
+#include <brpc/server.h>
+#include <gtest/gtest.h>
+
+#include <future>
+#include <thread>
+
+#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<ResultBlockBufferBase> buffer;
+        ASSERT_TRUE(_mgr.create_sender(_id, 16, &buffer, &_state, true, schema).ok());
+        _buffer = std::dynamic_pointer_cast<ArrowFlightResultBlockBuffer>(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<QueryStatement>(_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<Block>(ColumnHelper::create_block<DataTypeInt64>({1, 2}));
+        ASSERT_TRUE(_buffer->add_batch(&_state, block).ok());
+        ASSERT_FALSE(_dep->ready());
+    }
+    bool registered() {
+        std::shared_ptr<ArrowFlightResultBlockBuffer> found;
+        return _mgr.find_buffer(_id, found).ok();
+    }
+    MockRuntimeState _state;
+    ResultBufferMgr _mgr;
+    ResultBufferMgr* _previous_mgr = nullptr;
+    TUniqueId _id;
+    std::shared_ptr<ArrowFlightResultBlockBuffer> _buffer;
+    std::shared_ptr<Dependency> _dep;
+    std::shared_ptr<QueryStatement> _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> 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<ResultBlockBufferBase> 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<ArrowFlightResultBlockBuffer>(sibling);
+    auto block = std::make_shared<Block>(ColumnHelper::create_block<DataTypeInt64>({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<ArrowFlightResultBlockBuffer> 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<arrow::RecordBatch> batch;
+    ASSERT_TRUE(reader->ReadNext(&batch).ok());
+    EXPECT_EQ(batch, nullptr);
+    EXPECT_TRUE(reader->Close().ok());
+    std::shared_ptr<Block> 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<bool> cancelled = false;
+    std::promise<void> checked;
+    std::atomic<int> 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<arrow::RecordBatch> 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<Block>();
+        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<void> retried;
+    std::atomic<int> 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<arrow::RecordBatch> 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<Block>(ColumnHelper::create_block<DataTypeInt64>({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<Block>(ColumnHelper::create_block<DataTypeString>({"bad"}));
+    ASSERT_TRUE(_buffer->add_batch(&_state, block).ok());
+    auto result = ArrowFlightBatchLocalReader::Create(_statement);
+    ASSERT_TRUE(result.ok()) << result.status();
+    std::shared_ptr<arrow::RecordBatch> 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<LoadStreamMgr>(1);
+        _service = std::make_unique<PInternalService>(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<MockFragmentManager>(_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<PBackendService_Stub> _cache;
+    BrpcClientCache<PBackendService_Stub>* _previous_cache = nullptr;
+    std::unique_ptr<LoadStreamMgr> _previous_load_mgr;
+    std::unique_ptr<PInternalService> _service;
+    brpc::Server _server;
+    DNSCache* _previous_dns_cache = nullptr;
+    Status _cancel_status;
+    std::unique_ptr<MockFragmentManager> _fragment_mgr;
+    FragmentMgr* _previous_fragment_mgr = nullptr;
+};
+
+TEST_F(ArrowFlightRemoteCancelTest, CancellationInterruptsPendingBrpcFetch) {
+    std::atomic<bool> 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<arrow::RecordBatch> 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<ErrorCode::LIMIT_REACH>("limit reached"),
+                               Status::Error<ErrorCode::FINISHED>("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