Skip to content

[grpc-query] Respect client cancellations/disconnections#19005

Open
bsmithgall wants to merge 3 commits intoapache:masterfrom
bsmithgall:grpc-query/handle-cancelled
Open

[grpc-query] Respect client cancellations/disconnections#19005
bsmithgall wants to merge 3 commits intoapache:masterfrom
bsmithgall:grpc-query/handle-cancelled

Conversation

@bsmithgall
Copy link
Contributor

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:

  1. Registering a Context.CancellationListener
  2. Invoking DirectStatement.cancel() for SQL queries and QueryScheduler.cancelQuery() for native queries
  3. Mapping QueryInterruptedException to appropriate gRPC status codes

This change uses an AtomicReference to handle a race condition where the client cancels after the listener is registered but before the callback is set. We use getAndSet() 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
  • QueryService
  • QueryDriver
  • GrpcEndpointInitializer
  • QueryServiceTest

This PR has:

  • been self-reviewed.
  • a release note entry in the PR description.
  • added unit tests or modified existing tests to cover new code paths, ensuring the threshold for code coverage is met.

`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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant