liviazhu commented on code in PR #59144:
URL: https://github.com/apache/spark/pull/59144#discussion_r4147431771
##########
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:
Unfortunately, publishing the provider before obtaining the store is subject
to the following race condition:
1. Task A creates provider and publishes to loadedProviders
2. Task B gets loadedProviders lock, sees ownership of the provider has
moved to a different executor, removes it and adds to unloadedProvidersToClose
queue. Maintenance then closes it.
3. Task A tries to create a state store instance from the provider and runs
into the same `STATE_STORE_INVALID_STATE_MACHINE_TRANSITION` error
I think the "invariant" that "loadedProviders holds the only provider for an
id on the executor" is not true today anyway, since there can be multiple
providers for an id on an executor: one can be in loadedProviders, and other
can be waiting to be unloaded in `unloadedProvidersToClose` (and some can be
actively used by tasks). If the invariant should be "loadedProviders holds the
only provider for an id on the executor outside of the close queue", I don't
really see why it needs to be an invariant? They will not share the same native
RocksDB, and if they race uploading state files, it will be first-writer-wins
(either writing v1 files or reporting id to checkpoint v2), which is fine.
Let me know if you disagree!
--
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]
