[SPARK-43803] [SS] [CONNECT] Improve awaitTermination() to handle client disconnects

[pengzhon-db]

What changes were proposed in this pull request?

Streaming awaitTermination() is a long running API. Currently, it keeps running on server even if client disconnects. This change periodically checks if client has disconnected. If so, we can stop the operation and release resources.

We use gRPC Context.isCancelled() to determine if client has disconnected and response cannot be returned to the client. Here is the reference of of isCancelled():

Return whether this RPC failed before the server could provide its status back to the client.

This could be because of explicit API cancellation from the client-side or server-side, because of deadline exceeded, network connection reset, HTTP/2 parameter configuration (e.g., max message size, max connection age), etc. It does NOT include failure due to a non-OK status return from the server application's request handler, including Status::CANCELLED.

IsCancelled is always safe to call when using sync or callback API. When using async API, it is only safe to call IsCancelled after the AsyncNotifyWhenDone tag has been delivered. Thread-safe.

Why are the changes needed?

The change improves handling of awaitTermination(). It avoids resource waste of server side.

Does this PR introduce any user-facing change?

No

How was this patch tested?

Existing unit tests.
Manually tested on local with:

1. start spark connect
2. create a streaming query
3. call query.awaitTermination()
4. exit() the client to disconnect it
5. check that an error (RPC context is cancelled when executing awaitTermination()) is logged on server. It proves that awaitTermination() is exited on the server side when client disconnects

>>> query = (
...  spark
...  .readStream
...  .format("rate")
...  .option("numPartitions", "1")
...  .load()
...  .writeStream
...  .format("memory")
...  .queryName("tableName_35")
...  .start()
... )

>>>
>>> query.awaitTermination()
...
>>> exit()

[rangadi]

Left a few comments. Can we add unit test for this?

[pengzhon-db]

Left a few comments. Can we add unit test for this?

What kind of unit test are you referring to? We have existing Spark connect awaittermination unit test. Right now we don't have way to simulate client disconnect from python side. Do you think we should add some unit test for that?

[rangadi]

What kind of unit test are you referring to? We have existing Spark connect awaittermination unit test. Right now we don't have way to simulate client disconnect from python side. Do you think we should add some unit test for that?

We could disconnect in the test (since we have access to client). Lets may be that could be part of PR for client side changes.

[pengzhon-db]

@HyukjinKwon Could u review this PR?

[HyukjinKwon]

Is this good to go @rangadi ?

[rangadi]

@HyukjinKwon LGTM. will check CI failure.

[rangadi]

@pengzhon-db please fix scala-format error.

[pengzhon-db]

@HyukjinKwon can you help merge this? Thanks

[HyukjinKwon]

Merged to master.