(spark) branch branch-4.x updated: [SPARK-58844][UDF] Fix a Cancel-after-Finish flake in the UDF worker protocol test client

[yi...@apache.org]

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 <mo...@databricks.com>
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 <mo...@databricks.com>
    Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
    (cherry picked from commit de14ec8ba53532602a4645584e1673f8f11a09d2)
    Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
---
 .../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: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org