[Transport-gRPC] Wire gRPC call cancellation into search task cancellation

[pseudo-nymous]

Description

Cancels the server-side search task when a gRPC search call is cancelled (client cancels the RPC, disconnects, or the call deadline is exceeded), bringing the transport-grpc transport to parity with the HTTP transport.

Problem

Over the transport-grpc transport, a search whose client had already gone away kept running to completion server-side — we observed 17–30s of wasted execution after cancellation, consuming CPU, memory, and search-thread-pool capacity for a result no client would read. A client that times out and retries can leave multiple orphaned searches running, amplifying load when the cluster is already slow.

Observed impact (production)

First caught during a segment merge — gRPC search requests spiked to >10s, driving up p99. Merges are just one trigger; this hurts during any period of elevated search latency (GC pauses, high load, expensive queries, cold caches, etc.). Whenever a search is slow enough that the client gives up (cancel, disconnect, or deadline), the abandoned request kept running across all shards, compounding the slowdown that made the client leave.

Representative slow log — a lookup over 256 shards spending 17.3s almost entirely in the query phase (query=17318 of took_millis=17319), 0 hits:

took[17.3s], took_millis[17319], phase_took_millis[{expand=0, query=17318, fetch=0}], total_hits[0 hits], search_type[QUERY_THEN_FETCH], shards[{total:256, successful:256, skipped:0, failed:0}], source[{"from":0,"size":100,"query":{"bool":{"filter":[{"multi_match":{"query":"xyz","fields":["user_contactinfo.*^1.0"],"type":"best_fields","operator":"OR","analyzer":"phone_search_analyzer"}}],"must_not":[{"exists":{"field":"deleted_at"}}]}},"stored_fields":["uuid"],"sort":[{"email_address":{"order":"asc","missing":""}}]}], id[]

Since the time is spent almost entirely in the query phase, cancelling the coordinator task and its multi shard sub-tasks (what this change does) stops the work instead of letting it run the full 17.3s — relieving the p99 spike during merges and any other period of elevated latency.

Root cause

The HTTP path wraps the node client in RestCancellableNodeClient, which captures the search Task via NodeClient.executeLocally(...) and, on HttpChannel close, issues a CancelTasksRequest for it. The gRPC path (SearchServiceImpl.search) instead called client.search(searchRequest, listener), which discards the task, and never wired anything to gRPC's cancellation signal — so gRPC search had no cancellation-on-client-disconnect behavior at all.

Fix

Mirror the HTTP design with a dedicated collaborator so the service handler stays focused and cancellation stays encapsulated:

• GrpcCancellableNodeClient (new) — a FilterClient that:
  • executes the action via NodeClient.executeLocally(...) to obtain the Task;
  • registers ServerCallStreamObserver.setOnCancelHandler(...) to issue a CancelTasksRequest for that task (as the tasks-origin system user, exactly like RestCancellableNodeClient);
  • guards the register-time race with an isCancelled() check so a call cancelled between executeLocally and handler registration is still cancelled.
• SearchServiceImpl — routes the search through GrpcCancellableNodeClient when the call exposes a cancellation signal (the observer is a ServerCallStreamObserver and the client is a NodeClient), and otherwise falls back to the existing path unchanged.
• Testing in progress on real server.

Notes / scope

• No behavior change on the happy path: client.search(...) already routes through executeLocally internally (discarding the task), so this only retains the task handle to enable cancellation.
• No new public API, request field, or setting. Users don't opt in — abandoned work simply stops.
• Cancellation flows through the existing CancelTasks machinery, which propagates the ban to shard-level child tasks on data nodes, so the actual query/fetch work stops — not just the coordinator.
• Scope is limited to the search endpoint; PIT/other gRPC endpoints are unchanged.

Files changed

File	Change
.../transport/grpc/client/GrpcCancellableNodeClient.java	New: cancellation-aware FilterClient
.../transport/grpc/client/package-info.java	New: package docs
.../transport/grpc/services/SearchServiceImpl.java	Route search through the cancellable client when a cancellation signal is available
.../transport/grpc/client/GrpcCancellableNodeClientTests.java	New: unit tests
.../transport/grpc/services/SearchServiceImplTests.java	Added: delegation test
.../transport/grpc/SearchCancellationIT.java	New: integration test

Related Issues

Resolves #22428

Testing

• Unit — GrpcCancellableNodeClientTests: task is captured via executeLocally; firing the cancel handler cancels the captured task; an already-cancelled call cancels immediately; a normally-completing call leaves the task running.
• Unit — SearchServiceImplTests: with a ServerCallStreamObserver, the search is routed through the cancellable path (executeLocally + setOnCancelHandler, not client.search); existing tests unchanged.
• Integration — SearchCancellationIT: boots the real gRPC transport, holds a search in the query phase, cancels the gRPC call via a cancellable context, and asserts a CancelTasks is issued for the search task. Fully event-driven (latches, no assertBusy/sleeps). Verified it fails with the fix disabled and passes with it enabled.
• Manual, end-to-end (local, via grpcurl): on a node running the gRPC transport, ingest a large enough dataset that a search runs for several seconds (bulk-index many documents), then send an expensive gRPC search through grpcurl with a short call deadline (-max-time) so the client abandons the call while the search is still executing. With the fix, the server-side task is cancelled as soon as the deadline expires — confirmed via the DEBUG cancel log and the tasks API showing the search task cancelled=true and then unwinding — instead of running to natural completion. Without the fix, the same request keeps running server-side for the query's full duration after grpcurl has already returned.
• Real cluster (expensive query): the same approach applies on a live cluster — issue an expensive/slow gRPC search (e.g. a costly aggregation or regexp/wildcard query spanning many shards) so it runs long enough to observe, cancel the call (deadline, disconnect, or explicit cancel), and confirm via the tasks API that the search task and its shard sub-tasks are cancelled and stop promptly rather than running to completion.

Run (automated):

./gradlew :modules:transport-grpc:test \
  --tests "org.opensearch.transport.grpc.client.GrpcCancellableNodeClientTests" \
  --tests "org.opensearch.transport.grpc.services.SearchServiceImplTests"
./gradlew :modules:transport-grpc:internalClusterTest \
  --tests "org.opensearch.transport.grpc.SearchCancellationIT"

Manual reproduction (local, grpcurl):

# 1. Start a node with the gRPC auxiliary transport enabled (binds :9400)
./gradlew :run -Dtests.opensearch.aux.transport.types=transport-grpc

# 2. Bulk-index enough documents that a search takes several seconds to run
#    (repeat the bulk to grow the dataset until a query is reliably slow)

# 3. Fire an expensive search with a short call deadline so grpcurl gives up
#    while the search is still executing server-side
grpcurl -plaintext -max-time 2 \
  -d '{"index":["my-index"], ... expensive query ...}' \
  localhost:9400 org.opensearch.protobufs.services.SearchService/Search

# 4. Confirm the server cancelled the task instead of running to completion:
#    - the node log shows the DEBUG cancel line from GrpcCancellableNodeClient
#    - GET _tasks?actions=*search* shows the search task cancelled=true, then gone

Check List

• Functionality includes testing.
• API changes companion pull request created, if applicable. — N/A (no API/proto changes)
• Public documentation issue/PR created, if applicable. — N/A (no user-facing API/setting; internal behavior parity)

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Wrapped observer bypassed The search is dispatched with listener built from wrappedObserver (the circuit-breaker-wrapping observer), but searchClientFor inspects the raw responseObserver to decide whether to enable cancellation. If wrappedObserver is a ServerCallStreamObserver but the raw responseObserver is not (or vice versa in future refactors), the cancellation wiring will be inconsistent with what the listener actually writes to. More importantly, cancel handlers are only set on the raw observer; if the underlying gRPC framework only honors handlers set on the wrapped observer chain, cancellation may not fire. Verify ServerCallStreamObserver.setOnCancelHandler on the underlying observer is the correct hook here. private Client searchClientFor(StreamObserver<org.opensearch.protobufs.SearchResponse> responseObserver) { if (client instanceof NodeClient && responseObserver instanceof ServerCallStreamObserver) { return new GrpcCancellableNodeClient( (NodeClient) client, (ServerCallStreamObserver<org.opensearch.protobufs.SearchResponse>) responseObserver ); } return client; } Cancellation applies to all actions doExecute unconditionally uses executeLocally and registers a cancel handler for every action routed through this client. If this client is later reused for non-search actions or if the same serverCallObserver handles multiple executions, setOnCancelHandler will overwrite previously registered handlers (gRPC allows only one), potentially breaking cancellation semantics. Currently only one search call is dispatched per instance, so it works, but the class is easy to misuse. public <Request extends ActionRequest, Response extends ActionResponse> void doExecute( ActionType<Response> action, Request request, ActionListener<Response> listener ) { // Execute locally so we can obtain the Task and cancel it if the gRPC call is abandoned. Task task = client.executeLocally(action, request, listener); final TaskId taskId = new TaskId(client.getLocalNodeId(), task.getId()); serverCallObserver.setOnCancelHandler(() -> { logger.debug("gRPC call cancelled by client, cancelling task {}", taskId); cancelTask(taskId); }); // Guard against the race where the call is cancelled between executeLocally and setOnCancelHandler: // if it is already cancelled the handler may never fire, so cancel the task explicitly. if (serverCallObserver.isCancelled()) { logger.debug("gRPC call already cancelled, cancelling task {}", taskId); cancelTask(taskId); } }

[github-actions]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
General	Log task cancellation failures instead of swallowing The ActionListener.wrap(() -> {}) swallows any failure silently, which makes cancellation failures invisible and hard to diagnose. Log the failure at debug/warn level so operators can see when task cancellation itself fails. modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/client/GrpcCancellableNodeClient.java [87-91] protected void cancelTask(TaskId taskId) { CancelTasksRequest req = new CancelTasksRequest().setTaskId(taskId).setReason("gRPC client cancelled the request"); // force the origin to execute the cancellation as a system user - new OriginSettingClient(client, TASKS_ORIGIN).admin().cluster().cancelTasks(req, ActionListener.wrap(() -> {})); + new OriginSettingClient(client, TASKS_ORIGIN).admin() + .cluster() + .cancelTasks(req, ActionListener.wrap(r -> {}, e -> logger.debug("Failed to cancel task {}", taskId, e))); } Suggestion importance[1-10]: 5 __ Why: Logging cancellation failures improves observability and diagnostic ability, which is a valid minor improvement over silently swallowing errors.	Low
	Guard cancel handler registration against IllegalStateException ServerCallStreamObserver.setOnCancelHandler must be called before the response is sent or any messages are requested, and typically during call setup on the gRPC thread. Calling it here inside doExecute, which may be invoked after the call has already started sending, can throw IllegalStateException from gRPC. Consider registering the cancel handler before invoking the action, or catch/verify the observer state prior to registration. modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/client/GrpcCancellableNodeClient.java [65-78] -// Execute locally so we can obtain the Task and cancel it if the gRPC call is abandoned. +// Register the cancel handler before executing the action so gRPC accepts the handler +// registration (it must be set before responses are sent). Task task = client.executeLocally(action, request, listener); final TaskId taskId = new TaskId(client.getLocalNodeId(), task.getId()); -serverCallObserver.setOnCancelHandler(() -> { - logger.debug("gRPC call cancelled by client, cancelling task {}", taskId); - cancelTask(taskId); -}); +try { + serverCallObserver.setOnCancelHandler(() -> { + logger.debug("gRPC call cancelled by client, cancelling task {}", taskId); + cancelTask(taskId); + }); +} catch (IllegalStateException e) { + logger.debug("Unable to register cancel handler for task {}: {}", taskId, e.getMessage()); +} -// Guard against the race where the call is cancelled between executeLocally and setOnCancelHandler: -// if it is already cancelled the handler may never fire, so cancel the task explicitly. if (serverCallObserver.isCancelled()) { logger.debug("gRPC call already cancelled, cancelling task {}", taskId); cancelTask(taskId); } Suggestion importance[1-10]: 3 __ Why: The concern about setOnCancelHandler throwing IllegalStateException is somewhat speculative here since doExecute is called during the RPC handler before any response is sent. The suggestion adds defensive try/catch but does not clearly address a demonstrated bug.	Low

[github-actions]

✅ Gradle check result for a131ec6: SUCCESS

[codecov]

Codecov Report

❌ Patch coverage is 75.00000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 73.47%. Comparing base (27e5d58) to head (a131ec6).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
...rch/transport/grpc/services/SearchServiceImpl.java	75.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #22429      +/-   ##
============================================
+ Coverage     73.39%   73.47%   +0.07%     
- Complexity    76397    76503     +106     
============================================
  Files          6105     6105              
  Lines        346613   346616       +3     
  Branches      49888    49888              
============================================
+ Hits         254403   254664     +261     
+ Misses        71958    71689     -269     
- Partials      20252    20263      +11     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.