micheal-o commented on code in PR #58258: URL: https://github.com/apache/spark/pull/58258#discussion_r3868903428
########## sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreCoordinatorBatchReadSuite.scala: ########## @@ -0,0 +1,82 @@ +/* + * 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.execution.streaming.state + +import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.streaming.{OutputMode, StreamTest} + +/** + * Regression test for SPARK-58973: Stateful operations fail in non-streaming mode + * when StreamingQueryManager hasn't been initialized. + * + * This is a separate test file because: + * - StateDataSourceTestBase.beforeEach() calls spark.streams.stateStoreCoordinator which + * initializes the coordinator, masking the bug in all existing state data source tests. + * - We need private[state] access to stop the coordinator and verify coordinatorRef behavior. + */ +class StateStoreCoordinatorBatchReadSuite extends StreamTest { Review Comment: I don't think we should add a new test suite just for this single issue. There are existing test suite that covers these areas. See my other comment below. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala: ########## @@ -1174,17 +1174,30 @@ object StateStoreProvider extends Logging { private[state] def coordinatorRef: Option[StateStoreCoordinatorRef] = synchronized { val env = SparkEnv.get if (env != null) { - val isDriver = SparkContext.isDriver(env.executorId) - // If running locally, then the coordinator reference in stateStoreCoordinatorRef may have - // become inactive as SparkContext + SparkEnv may have been restarted. Hence, when running in - // driver, always recreate the reference. - if (isDriver || stateStoreCoordinatorRef == null) { - logDebug("Getting StateStoreCoordinatorRef") - stateStoreCoordinatorRef = StateStoreCoordinatorRef.forExecutor(env) + try { + val isDriver = SparkContext.isDriver(env.executorId) + // If running locally, then the coordinator reference in stateStoreCoordinatorRef may have + // become inactive as SparkContext + SparkEnv may have been restarted. Hence, when running + // in driver, always recreate the reference. + if (isDriver || stateStoreCoordinatorRef == null) { + logDebug("Getting StateStoreCoordinatorRef") + stateStoreCoordinatorRef = StateStoreCoordinatorRef.forExecutor(env) + } + logInfo(log"Retrieved reference to StateStoreCoordinator: " + + log"${MDC(LogKeys.STATE_STORE_COORDINATOR, stateStoreCoordinatorRef)}") + Some(stateStoreCoordinatorRef) + } catch { + // SPARK-58973: The StateStoreCoordinator endpoint is only registered when + // StreamingQueryManager is initialized (lazy val in SessionState since SPARK-29423). + // In non-streaming batch reads of state stores (e.g., spark.read.format("statestore")), + // the coordinator may not exist. Return None to make snapshot upload reporting + // best-effort rather than failing the read. + case e: SparkException => Review Comment: `SparkException` is a broad exception type, can we catch something more specific here? ########## sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreCoordinatorBatchReadSuite.scala: ########## @@ -0,0 +1,82 @@ +/* + * 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.execution.streaming.state + +import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.streaming.{OutputMode, StreamTest} + +/** + * Regression test for SPARK-58973: Stateful operations fail in non-streaming mode + * when StreamingQueryManager hasn't been initialized. + * + * This is a separate test file because: + * - StateDataSourceTestBase.beforeEach() calls spark.streams.stateStoreCoordinator which + * initializes the coordinator, masking the bug in all existing state data source tests. + * - We need private[state] access to stop the coordinator and verify coordinatorRef behavior. + */ +class StateStoreCoordinatorBatchReadSuite extends StreamTest { + import testImplicits._ + + test("SPARK-58973: coordinatorRef returns None when coordinator is not registered") { Review Comment: Can this be added to `StateStoreCoordinatorSuite` instead? ########## sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreCoordinatorBatchReadSuite.scala: ########## @@ -0,0 +1,82 @@ +/* + * 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.execution.streaming.state + +import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.streaming.{OutputMode, StreamTest} + +/** + * Regression test for SPARK-58973: Stateful operations fail in non-streaming mode + * when StreamingQueryManager hasn't been initialized. + * + * This is a separate test file because: + * - StateDataSourceTestBase.beforeEach() calls spark.streams.stateStoreCoordinator which + * initializes the coordinator, masking the bug in all existing state data source tests. + * - We need private[state] access to stop the coordinator and verify coordinatorRef behavior. + */ +class StateStoreCoordinatorBatchReadSuite extends StreamTest { + import testImplicits._ + + test("SPARK-58973: coordinatorRef returns None when coordinator is not registered") { + // Stop the coordinator to simulate a fresh session where + // StreamingQueryManager was never initialized (lazy val in SessionState). + spark.sessionState.streamingQueryManager.stateStoreCoordinator.stop() + StateStore.stop() + + // Unit-level verification: coordinatorRef should gracefully return None + // instead of throwing RpcEndpointNotFoundException. + assert(StateStoreProvider.coordinatorRef.isEmpty, + "coordinatorRef should return None when the coordinator endpoint is not registered") + } + + test("SPARK-58973: reading state store succeeds when coordinator is not initialized") { Review Comment: Can this be added to `StateDataSourceReadSuite` instead? -- 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]
