[grpc-query] Respect client cancellations/disconnections#19005
Conversation
`maven-enforce-plugin`'s `RequireUpperBoundDeps` rule caught this change due to a transitive dependency from `assertj-core:3.27.7`.
This change allows in-flight queries to be cancelled (for example, by a gRPC client) by accepting a disconnection callback.
When a gRPC client disconnects or cancels, forward that cancellation signal into the in-flight Druid query and return `Status.CANCELLED` by: 1. Registering a gRPC `Context.CancellationListener` 2. Invoking `DirectStatement.cancel()` for SQL queries and `QueryScheduler.cancelQuery()` for native queries 3. Mapping exceptions back to the correct status codes In addition, guard against possible cancellation in the event that the cancellation happens before the listener callback is fully registered.
There was a problem hiding this comment.
Pull request overview
This PR adds support for in-flight Druid query cancellation when gRPC clients disconnect or cancel their requests. Previously, queries would continue to run even after client disconnection. The implementation registers a gRPC Context.CancellationListener that forwards cancellation signals to Druid by invoking DirectStatement.cancel() for SQL queries and QueryScheduler.cancelQuery() for native queries, and maps QueryInterruptedException to appropriate gRPC status codes.
Changes:
- Added gRPC context cancellation listener and callback mechanism to properly handle client disconnections/cancellations
- Extended QueryDriver and QueryService to support query cancellation via callback pattern
- Updated exception handling to map QueryInterruptedException to CANCELLED or INTERNAL gRPC status codes
- Added QueryScheduler dependency to enable native query cancellation
- Upgraded bytebuddy from 1.17.7 to 1.18.3
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| pom.xml | Upgraded bytebuddy dependency version |
| QueryService.java | Added cancellation listener registration and QueryInterruptedException handling |
| QueryDriver.java | Added cancellation callback support for both SQL and native queries |
| GrpcEndpointInitializer.java | Added QueryScheduler injection to enable native query cancellation |
| QueryServiceTest.java | Added comprehensive tests for cancellation scenarios |
| TestServer.java, GrpcQueryTest.java, DriverTest.java, BasicAuthTest.java | Updated to pass QueryScheduler parameter |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if (queryScheduler != null) { | ||
| final String queryId = query.getId(); | ||
| cancelCallback.set(() -> queryScheduler.cancelQuery(queryId)); | ||
| } |
There was a problem hiding this comment.
The cancelCallback is set after query parsing and lifecycle initialization but before the query is executed (line 173). If the client cancels during the parse/initialize phase (lines 154-169), the cancellation will not be effective until after those operations complete. This could be significant for expensive parsing operations or complex query initialization.
Consider setting up the cancelCallback earlier, perhaps with a mechanism to cancel the QueryLifecycle itself during initialization, or at minimum add a cancellation check after initialization and before execution.
| Status.INTERNAL.withDescription(e.getMessage()) | ||
| )); | ||
| } | ||
| } |
There was a problem hiding this comment.
The exception handling in QueryService doesn't have a catch-all for other types of exceptions. If any exception other than ForbiddenException or QueryInterruptedException is thrown (e.g., RuntimeException from QueryDriver), it will propagate uncaught and gRPC will handle it with a generic INTERNAL error. This is inconsistent with typical error handling patterns and could result in poor error messages being sent to the client.
Consider adding a catch-all exception handler that maps to an appropriate gRPC status code.
| } | |
| } | |
| catch (Exception e) { | |
| responseObserver.onError(new StatusRuntimeException( | |
| Status.INTERNAL.withDescription(e.getMessage()) | |
| )); | |
| } |
| if (Context.current().isCancelled()) { | ||
| cancelOnce.run(); | ||
| } |
There was a problem hiding this comment.
The finally block's race condition handling at lines 86-88 has a flaw. If the context is cancelled after the query completes normally but before the finally block executes, this will invoke the cancellation callback even though the query already finished successfully. This could potentially cause issues if the cancellation callback tries to cancel an already-completed query.
A better approach would be to track whether the query has completed and only invoke the cancellation callback if it hasn't.
| catch (QueryInterruptedException e) { | ||
| throw e; | ||
| } |
There was a problem hiding this comment.
For native queries, the QueryLifecycle is not properly cleaned up when a QueryInterruptedException occurs. Unlike SQL queries (lines 257-261) which call stmt.reporter().failed(e) and stmt.close(), native queries simply re-throw the exception without cleanup.
The queryLifecycle object should have emitLogsAndMetrics() called (as noted in the TODO comment on line 193) and any other necessary cleanup should be performed. This could lead to resource leaks and incomplete metrics/logging for cancelled native queries.
gianm
left a comment
gianm
left a comment
There was a problem hiding this comment.
Took a quick look. Generally looks OK, except for the TODO comment.
Description
Add support for in-flight Druid query cancellation after a gRPC client disconnects or cancels. Previously, these events were ignored, leading to queries continuing to run. Now, these cancellations forward that signal into Druid by:
Context.CancellationListenerDirectStatement.cancel()for SQL queries andQueryScheduler.cancelQuery()for native queriesQueryInterruptedExceptionto appropriate gRPC status codesThis change uses an
AtomicReferenceto handle a race condition where the client cancels after the listener is registered but before the callback is set. We usegetAndSet()to ensure cancellation is fired at most once.Release note
The gRPC query extension now cancels in-flight queries when clients cancel/disconnect.
Key changed/added classes in this PR
QueryServiceQueryDriverGrpcEndpointInitializerQueryServiceTestThis PR has: