HyukjinKwon commented on code in PR #59144:
URL: https://github.com/apache/spark/pull/59144#discussion_r4140900374
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala:
##########
@@ -1625,13 +1628,69 @@ object StateStore extends Logging {
if (version < 0) {
throw QueryExecutionErrors.unexpectedStateStoreVersion(version)
}
- val (storeProvider, shouldForceSnapshotUpload) =
getStateStoreProvider(storeProviderId,
- keySchema, valueSchema, keyStateEncoderSpec, useColumnFamilies,
storeConf, hadoopConf,
- useMultipleValuesPerKey, stateSchemaBroadcast)
- storeProvider.getStore(version, stateStoreCkptId,
shouldForceSnapshotUpload)
+ withClosedProviderRetry(storeProviderId) { createUnpublishedProvider =>
+ getStateStoreProvider(storeProviderId,
+ keySchema, valueSchema, keyStateEncoderSpec, useColumnFamilies,
storeConf, hadoopConf,
+ useMultipleValuesPerKey, stateSchemaBroadcast,
createUnpublishedProvider)
+ } { (storeProvider, shouldForceSnapshotUpload) =>
+ storeProvider.getStore(version, stateStoreCkptId,
shouldForceSnapshotUpload)
+ }
}
// scalastyle:on
+ private def withClosedProviderRetry[T](
+ storeProviderId: StateStoreProviderId)(
+ getProvider: Boolean => (StateStoreProvider, Boolean))(
+ loadStore: (StateStoreProvider, Boolean) => T): T = {
+ val (firstProvider, firstShouldForceSnapshotUpload) = getProvider(false)
+ try {
+ loadStore(firstProvider, firstShouldForceSnapshotUpload)
+ } catch {
+ case error: StateStoreInvalidStateMachineTransition
+ if isClosedProviderLoadTransition(error) =>
+ loadedProviders.synchronized {
+ if (loadedProviders.get(storeProviderId).exists(_ eq firstProvider))
{
+ loadedProviders.remove(storeProviderId)
+ } else {
+ None
+ }
+ }.foreach { provider =>
+ unloadedProvidersToClose.add(
+ (storeProviderId, provider, MaintenanceOpRequest.All))
+ }
+ // Create the replacement without publishing it in loadedProviders yet.
+ val (replacementProvider, replacementShouldForceSnapshotUpload) =
getProvider(true)
Review Comment:
**Non-blocking (P2):** This replacement is created with
`createUnpublishedProvider = true`, so it is not in `loadedProviders` while
`loadStore` runs. If another task on the same executor calls `StateStore.get`
for the same `StateStoreProviderId` during that window (the Cartesian siblings
this PR describes are exactly that case), `getOrElseUpdate` creates and
publishes a second provider. That task then acquires its store immediately,
because the two providers have independent `RocksDBStateMachine`s and separate
local RocksDB directories. In the baseline, both callers would share the single
cached provider and the second would wait in `acquireStamp`.
When this caller then finds `loadedProviders.contains(storeProviderId)`, it
queues `replacementProvider` for close but still returns the store it loaded
from it. Two stores for the same partition and version are now acquired
concurrently on one executor, and both can commit checkpoint files for that
version. Maintenance also runs snapshot/cleanup against the replacement and
blocks in `close()` until release (up to `lockAcquireTimeoutMs`).
Could the retry instead go through the cache? For example, remove the
identical stale entry under `loadedProviders.synchronized` as you do now, then
call the normal `getStateStoreProvider` path (`getOrElseUpdate`) so the
replacement is published immediately and concurrent callers share it. That
would drop `createUnpublishedProvider` and the post-load publish/loser branch
entirely.
See **Shared repair plan 1** in the review body.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreLockHardeningSuite.scala:
##########
@@ -250,13 +286,86 @@ class RocksDBStateStoreLockHardeningSuite extends
SparkFunSuite
val stateMachineObj = provider invokePrivate stateMachine()
stateMachineObj.asInstanceOf[RocksDBStateMachine].close()
- // Attempting to get a new store after close should fail
- // with StateStoreInvalidStateMachineTransition
val exception = intercept[StateStoreInvalidStateMachineTransition] {
provider.getStore(1)
}
- assert(exception.getMessage.contains("Old state: CLOSED"))
+ checkError(
+ exception,
+ condition = "STATE_STORE_INVALID_STATE_MACHINE_TRANSITION",
+ parameters = Map(
+ "oldState" -> "CLOSED",
+ "newState" -> "ACQUIRED",
+ "operation" -> "load",
+ "storeId" -> provider.stateStoreId.toString))
+ }
+
+ test("SPARK-59877: StateStore retries a cached closed provider") {
Review Comment:
**Nit (P3):** This test covers `StateStore.get` with a replacement that
loads successfully and wins the publish. None of the new tests call
`getReadOnly` on a closed cached provider, issue a concurrent same-id
`StateStore.get` while the replacement is blocked, or make the replacement's
`getStore` throw. So regressions in read-only routing, the losing-publish
branch, or the failure cleanup (which would leak a RocksDB instance) would not
be caught. `BlockingReplacementProvider` already makes these interleavings
deterministic, so a concurrent second caller plus a throwing variant should be
cheap to add.
See **Shared repair plan 1** in the review body.
--
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]
