spark-test commented on code in PR #59549: URL: https://github.com/apache/spark/pull/59549#discussion_r4422487096
########## sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectSessionValidationSuite.scala: ########## @@ -0,0 +1,149 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connect.client + +import java.util.UUID +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicInteger + +import scala.concurrent.duration._ + +import io.grpc.{Status, StatusRuntimeException} +import io.grpc.netty.NettyServerBuilder +import io.grpc.stub.StreamObserver +import org.scalatest.concurrent.{Signaler, ThreadSignaler, TimeLimits} + +import org.apache.spark.connect.proto +import org.apache.spark.sql.connect.test.ConnectFunSuite + +class SparkConnectSessionValidationSuite extends ConnectFunSuite with TimeLimits { + + private implicit val signaler: Signaler = ThreadSignaler + + private class ValidationService + extends proto.SparkConnectServiceGrpc.SparkConnectServiceImplBase { + val calls = new AtomicInteger() + @volatile var failure: Option[Status] = None + @volatile var respond = true + @volatile var lastRequest: proto.ConfigRequest = _ + @volatile var serverSessionId = UUID.randomUUID().toString + + override def config( + request: proto.ConfigRequest, + observer: StreamObserver[proto.ConfigResponse]): Unit = { + lastRequest = request + calls.incrementAndGet() + if (respond) { + failure match { + case Some(status) => + // Fail only once so unexpected retries would succeed. + failure = None + observer.onError(status.asRuntimeException()) + case None => + observer.onNext( + proto.ConfigResponse + .newBuilder() + .setSessionId(request.getSessionId) + .setServerSideSessionId(serverSessionId) + .build()) + observer.onCompleted() + } + } + } + } + + private def withClient(f: (SparkConnectClient, ValidationService) => Unit): Unit = { + val service = new ValidationService + val server = NettyServerBuilder.forPort(0).addService(service).build().start() + val client = SparkConnectClient.builder().port(server.getPort).build() + try { + f(client, service) + } finally { + client.shutdown() + server.shutdownNow() + assert(server.awaitTermination(5, TimeUnit.SECONDS)) + } + } + + test("validateSession checks the server, including when the timeout is zero") { + withClient { (client, service) => + assert(client.validateSession(0)) + assert(client.validateSession(5)) + assert(service.calls.get() == 2) + assert(service.lastRequest.getSessionId == client.sessionId) + assert(service.lastRequest.getClientObservedServerSideSessionId == service.serverSessionId) + } + } + + test("validateSession rejects a negative timeout without contacting the server") { + withClient { (client, service) => + intercept[IllegalArgumentException] { + client.validateSession(-1) + } + assert(service.calls.get() == 0) + } + } + + test("validateSession returns false when the server closes the session") { + withClient { (client, service) => + assert(client.validateSession(5)) + service.failure = Some( + Status.INTERNAL.withDescription( + "[INVALID_HANDLE.SESSION_CLOSED] Session was closed. SQLSTATE: 08003")) + assert(!client.validateSession(5)) + } + } + + test("validateSession returns false when the server-side session changes") { + withClient { (client, service) => + assert(client.validateSession(5)) + service.serverSessionId = UUID.randomUUID().toString + assert(!client.validateSession(5)) Review Comment: **Nit (P3):** `validateSession` documents, and implements at `SparkConnectClient.scala:199-201`, a false result without a server request once the local latch is false. No new test validates again after the latch flips. This test ends at the call that first sets it, and the JDBC suite never sets it, because `SESSION_CLOSED` does not latch. Deleting the `if (!isSessionValid) return false` guard therefore leaves every test green. Each `isValid` call on a session already known to be dead would then make an extra round trip, which can block when the timeout is 0. One more call here pins the guard: ```suggestion assert(!client.validateSession(5)) // The session is now locally invalid, so validation must not contact the server again. assert(!client.validateSession(5)) assert(service.calls.get() == 2) ``` **Verification:** - **Regression:** After the server-side session ID changes and validation first returns false, a further validation returns false without the server receiving another request. ########## sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala: ########## @@ -129,6 +129,18 @@ private[connect] class CustomSparkConnectBlockingStub( } } + def validateSession(request: ConfigRequest, timeout: Int): Unit = { + val validationStub = if (timeout > 0) { + stub.withDeadlineAfter(timeout.toLong, TimeUnit.SECONDS) + } else { + stub Review Comment: **Non-blocking (P2):** For `timeout == 0` this branch sends the validation Config RPC on the bare stub, so `isValid(0)` waits for the server with no deadline at all. The regular `config` path bounds the same RPC with `withDeadline(stubState.rpcDeadlines.config)` (10 minutes by default). The `RpcDeadlines` contract leaves only non-reattachable `ExecutePlan` and `ReleaseExecute` without a deadline. JDBC's "0 means no timeout" removes the validation-specific timeout, not this client-level bound. If the server accepts the call but never completes it, the caller blocks forever, because keepalive only detects dead transports. That caller is typically a pool thread validating with `isValid(0)`. Before this PR, `isValid` answered at once from local state. Positive timeouts are unaffected. Falling back to the configured Config deadline keeps the single-attempt behavior: ```suggestion withDeadline(stubState.rpcDeadlines.config) ``` The `validateSession` Scaladoc in `SparkConnectClient` ("Zero means no deadline") would then say that zero applies no validation-specific timeout, while the client's Config RPC deadline still bounds the call. A `jdbc:sc://` URL cannot set `rpcDeadlines`, so JDBC connections get the 10-minute default. **Verification:** - **Regression:** With a client whose Config RPC deadline is configured to a short value and a server that accepts the validation request but never answers, validating with timeout 0 returns false within a bounded time instead of hanging. - **Behavior:** Validating with timeout 0 against a responsive server still succeeds, and validating after the server has closed the session still returns false. - **Compatibility:** Validating with a positive timeout is still bounded by the caller's timeout rather than the configured Config RPC deadline, and each validation still sends exactly one request without retries. ########## sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectSessionValidationSuite.scala: ########## @@ -0,0 +1,149 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connect.client + +import java.util.UUID +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicInteger + +import scala.concurrent.duration._ + +import io.grpc.{Status, StatusRuntimeException} +import io.grpc.netty.NettyServerBuilder +import io.grpc.stub.StreamObserver +import org.scalatest.concurrent.{Signaler, ThreadSignaler, TimeLimits} + +import org.apache.spark.connect.proto +import org.apache.spark.sql.connect.test.ConnectFunSuite + +class SparkConnectSessionValidationSuite extends ConnectFunSuite with TimeLimits { + + private implicit val signaler: Signaler = ThreadSignaler + + private class ValidationService + extends proto.SparkConnectServiceGrpc.SparkConnectServiceImplBase { + val calls = new AtomicInteger() + @volatile var failure: Option[Status] = None + @volatile var respond = true + @volatile var lastRequest: proto.ConfigRequest = _ + @volatile var serverSessionId = UUID.randomUUID().toString + + override def config( + request: proto.ConfigRequest, + observer: StreamObserver[proto.ConfigResponse]): Unit = { + lastRequest = request + calls.incrementAndGet() + if (respond) { + failure match { + case Some(status) => + // Fail only once so unexpected retries would succeed. + failure = None + observer.onError(status.asRuntimeException()) + case None => + observer.onNext( + proto.ConfigResponse + .newBuilder() + .setSessionId(request.getSessionId) + .setServerSideSessionId(serverSessionId) + .build()) + observer.onCompleted() + } + } + } + } + + private def withClient(f: (SparkConnectClient, ValidationService) => Unit): Unit = { + val service = new ValidationService + val server = NettyServerBuilder.forPort(0).addService(service).build().start() + val client = SparkConnectClient.builder().port(server.getPort).build() + try { + f(client, service) + } finally { + client.shutdown() + server.shutdownNow() + assert(server.awaitTermination(5, TimeUnit.SECONDS)) + } + } + + test("validateSession checks the server, including when the timeout is zero") { + withClient { (client, service) => + assert(client.validateSession(0)) + assert(client.validateSession(5)) + assert(service.calls.get() == 2) + assert(service.lastRequest.getSessionId == client.sessionId) + assert(service.lastRequest.getClientObservedServerSideSessionId == service.serverSessionId) + } + } + + test("validateSession rejects a negative timeout without contacting the server") { + withClient { (client, service) => + intercept[IllegalArgumentException] { + client.validateSession(-1) + } + assert(service.calls.get() == 0) + } + } + + test("validateSession returns false when the server closes the session") { + withClient { (client, service) => + assert(client.validateSession(5)) + service.failure = Some( + Status.INTERNAL.withDescription( + "[INVALID_HANDLE.SESSION_CLOSED] Session was closed. SQLSTATE: 08003")) + assert(!client.validateSession(5)) + } + } + + test("validateSession returns false when the server-side session changes") { + withClient { (client, service) => + assert(client.validateSession(5)) + service.serverSessionId = UUID.randomUUID().toString + assert(!client.validateSession(5)) + } + } + + test("validation stub exceeds its deadline when the server does not respond") { + withClient { (client, service) => + assert(client.validateSession(0)) + val stub = new CustomSparkConnectBlockingStub( + client.channel, + new SparkConnectStubState(client.channel, client.configuration)) + service.respond = false + failAfter(60.seconds) { + val e = intercept[StatusRuntimeException] { + stub.validateSession(service.lastRequest, 1) Review Comment: **Nit (P3):** This is the only test against a server that never answers, and it calls a separately built stub directly. Every `client.validateSession(n)` call and every JDBC `isValid(n)` call in the new suites goes to a server that answers immediately. The `DEADLINE_EXCEEDED` case below is a status the mock returns at once. Replacing the forwarded timeout with 0 in `SparkConnectClient.validateSession` or `SparkConnectConnection.isValid` keeps every test green. Yet `isValid(5)` against an unresponsive server would then wait indefinitely, or, after the zero-timeout fallback lands, for up to the 10-minute Config RPC deadline. The existing `ValidationService` can drive the client entry point directly, for example: ```scala test("validateSession honors a positive timeout when the server does not respond") { withClient { (client, service) => assert(client.validateSession(5)) service.respond = false failAfter(60.seconds) { assert(!client.validateSession(1)) } } } ``` A JDBC connection opened on such a mock with `jdbc:sc://localhost:<port>` covers `isValid(1)` the same way. **Verification:** - **Regression:** Against a server that accepts the validation request but never answers, client-level validation with a short positive timeout returns false within a bounded time on a client whose Config RPC deadline is far longer than that bound, and JDBC Connection.isValid with a short positive timeout does the same. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
