This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new af3df8f2f6fc [SPARK-58844][UDF] Fix a Cancel-after-Finish flake in the
UDF worker protocol test client
af3df8f2f6fc is described below
commit af3df8f2f6fc34d394dbf8c3bb6ac4de68f6e780
Author: Mohamed Gallab <[email protected]>
AuthorDate: Wed Aug 19 20:13:11 2026 +0000
[SPARK-58844][UDF] Fix a Cancel-after-Finish flake in the UDF worker
protocol test client
### What changes were proposed in this pull request?
`EchoProtocolSuite`'s `EngineClient` test helper touches its request
`StreamObserver` from two threads without mutual exclusion. `sendCancel` does a
check-then-act:
- test thread: `if (requestCompleted.get()) return` reads `false`, then
calls `requestObserver.onNext(Cancel)`;
- meanwhile the response-observer callback thread runs
`completeRequestStream()` (`requestObserver.onCompleted()`), half-closing the
request stream.
If the half-close lands between the `get()` and the `onNext`, the `onNext`
throws, because gRPC does not permit concurrent calls on a `StreamObserver`.
This adds a `requestLock` — mirroring the existing `responseLock`, which
already serializes writes to the *response* observer — and guards both
`sendCancel` (check + send) and `completeRequestStream` (CAS + half-close) so
they can no longer interleave. Test-only change.
### Why are the changes needed?
`EchoProtocolSuite`'s `cancel: engine sends Cancel after Finish` case sends
a trailing `Cancel` exactly as the `FinishResponse` terminator arrives, so it
hits the race above and intermittently fails with `IllegalStateException: call
was half-closed` (or `Stream is already completed, no further calls are
allowed`). It passes in isolation and only fails under scheduling load, so it
can fail CI unpredictably.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing `EchoProtocolSuite` cases. The race window is tiny, so to confirm
the cause deterministically I temporarily widened it (a `Thread.sleep` between
the `get()` and the `onNext` in `sendCancel`): with the widening the `cancel:
engine sends Cancel after Finish` case failed on every run, and with
`requestLock` in place it passed on every run with the same widening —
confirming the lock closes the exact window. The temporary widening was then
removed.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Anthropic)
Closes #58079 from
mohamed-gallab-db/minor-udf-worker-cancel-after-finish-flake.
Authored-by: Mohamed Gallab <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
(cherry picked from commit de14ec8ba53532602a4645584e1673f8f11a09d2)
Signed-off-by: Yicong-Huang <[email protected]>
---
.../spark/udf/worker/grpc/EchoProtocolSuite.scala | 27 +++++++++++++---------
1 file changed, 16 insertions(+), 11 deletions(-)
diff --git
a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/EchoProtocolSuite.scala
b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/EchoProtocolSuite.scala
index eae6ef1a639a..4de834220839 100644
---
a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/EchoProtocolSuite.scala
+++
b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/EchoProtocolSuite.scala
@@ -469,6 +469,10 @@ class EchoProtocolSuite extends AnyFunSuite with
BeforeAndAfterEach {
@volatile var executionError: Option[ExecutionError] = None
@volatile var streamError: Option[Throwable] = None
private val requestCompleted = new AtomicBoolean(false)
+ // gRPC forbids concurrent calls to the request StreamObserver too: the
response
+ // observer half-closes it (completeRequestStream) on a callback thread
while the
+ // test thread may still be sending a trailing Cancel. Serialize both
through this.
+ private val requestLock = new Object
// Counted down on InitResponse (success or failure) or on terminal error.
// The engine MUST wait for this before sending any DataRequest or Finish.
private val initResponseLatch = new CountDownLatch(1)
@@ -637,22 +641,23 @@ class EchoProtocolSuite extends AnyFunSuite with
BeforeAndAfterEach {
}
}
- def sendCancel(reason: String = ""): Unit = {
- // If a terminator already arrived (FinishResponse / CancelResponse),
- // the request stream has been half-closed and Cancel arrives too
- // late -- silently ignore, matching the proto's Cancel-after-Finish
- // contract.
- if (requestCompleted.get()) return
- requestObserver.onNext(UdfRequest.newBuilder()
- .setControl(UdfControlRequest.newBuilder()
- .setCancel(Cancel.newBuilder().setReason(reason).build())
+ def sendCancel(reason: String = ""): Unit = requestLock.synchronized {
+ // If a terminator already arrived (FinishResponse / CancelResponse),
the request
+ // stream has been half-closed and Cancel arrives too late -- silently
ignore,
+ // matching the proto's Cancel-after-Finish contract. requestLock makes
this check
+ // and the send atomic against completeRequestStream's half-close.
+ if (!requestCompleted.get()) {
+ requestObserver.onNext(UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder()
+ .setCancel(Cancel.newBuilder().setReason(reason).build())
+ .build())
.build())
- .build())
+ }
// Request stream stays open until the response terminator arrives;
// completeRequestStream() is called by the response observer.
}
- def completeRequestStream(): Unit = {
+ def completeRequestStream(): Unit = requestLock.synchronized {
if (requestCompleted.compareAndSet(false, true)) {
requestObserver.onCompleted()
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
