morningman opened a new pull request, #69945:
URL: https://github.com/apache/doris/pull/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.
### Release note
Arrow Flight SQL: cancelling or closing a result stream before its end now
cancels the query on the BE that holds the stream's results and frees its
unread results there right away, instead of keeping them for minutes. A stream
aborted this way cannot be read again with the same ticket.
### Check List (For Author)
- Test: Unit Test / Manual test
- BE UT: `ArrowFlightCancelTest` (9, new), `ArrowFlightRemoteCancelTest`
(3, new), `ResultBufferMgrTest` (6), `ArrowResultBlockBufferTest` (4): pass 10
times in a row (ASAN, macOS); 10 mutations each caught
- `build-support/check-build-hygiene.sh` passes; clang-format 16 clean;
clang-tidy reports nothing on the changed lines except
`readability-inconsistent-ifelse-braces` inside the `DORIS_CHECK` macro
expansion
- Manual test on a local single-BE cluster: branch-4.x's
`test_flight_cancel_cleanup` fails before and passes after;
`arrow_flight_sql_p0` (20 suites) and `external_table_p0/remote_doris` (14
suites) pass
- No regression test in this PR: `test_flight_cancel_cleanup` needs the
FE route of the follow-up to pass on a multi-BE cluster, so it comes with it
- Behavior changed: Yes. A Flight result stream served by the BE that holds
its results and cancelled, closed or destroyed before EOF now cancels its query
on that BE and frees the query's Flight buffers there; a later `DoGet` with the
same ticket fails instead of resuming. On a stream served by another BE, the
abort only interrupts the pending fetch.
- Does this need documentation: No
Original author: @Gabriel39
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
