[SPARK-51596][SS] Fix concurrent StateStoreProvider maintenance and closing #50595

ericm-db · 2025-04-15T22:14:53Z

What changes were proposed in this pull request?

Moves the unload operation away from task thread into the maintenance thread. To ensure unloading still occurs ASAP (rather than potentially waiting for the maintenance interval) as was introduced by https://issues.apache.org/jira/browse/SPARK-33827, we immediately trigger a maintenance thread to do the unload.

This gives us an extra benefit that unloading other providers doesn't block the task thread. To capitalize on this, unload() should not hold the loadedProviders lock the entire time (which will block other task threads), but instead release it once it has deleted the unloading providers from the map and close the providers without the lock held.

Why are the changes needed?

Currently, both the task thread and maintenance thread can call unload() on a provider. This leads to a race condition where the maintenance could be conducting maintenance while the task thread is closing the provider, leading to unexpected behavior.

Does this PR introduce any user-facing change?

No

How was this patch tested?

Added unit test

Was this patch authored or co-authored using generative AI tooling?

No

anishshri-db · 2025-04-23T06:12:17Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

@@ -1157,6 +1160,69 @@ object StateStore extends Logging {
    }
  }

+  private def doMaintenanceOnProvider(id: StateStoreProviderId, provider: StateStoreProvider,


nit: lets move the args to new lines ?

anishshri-db · 2025-04-23T06:16:54Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

@@ -1009,14 +1013,45 @@ object StateStore extends Logging {

      val otherProviderIds = loadedProviders.keys.filter(_ != storeProviderId).toSeq
      val providerIdsToUnload = reportActiveStoreInstance(storeProviderId, otherProviderIds)
-      providerIdsToUnload.foreach(unload(_))
+      providerIdsToUnload.foreach(id => {
+        loadedProviders.remove(id).foreach( provider => {


We don't need to synchronize on loadedProviders here ?

Already in the synchronized block on line 922.

liviazhu-db · 2025-04-23T00:36:18Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+          // Trigger maintenance thread to immediately do maintenance on and close the provider.
+          // Doing maintenance first allows us to do maintenance for a constantly-moving state
+          // store.
+          logInfo(log"Task thread trigger maintenance on " +


Can you add some more info in this log line, like task ID? Also add that provider was removed from loadedProviders.

+1, that would help with investigation

include trigger maintenance to close provider

liviazhu-db · 2025-04-23T00:40:46Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

@@ -1157,6 +1160,69 @@ object StateStore extends Logging {
    }
  }

+  private def doMaintenanceOnProvider(id: StateStoreProviderId, provider: StateStoreProvider,


Can you add a comment saying that if alreadyRemovedFromLoadedProviders is false, we must already have the lock to process this partition?

liviazhu-db · 2025-04-24T16:25:51Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

@@ -1009,14 +1013,45 @@ object StateStore extends Logging {

      val otherProviderIds = loadedProviders.keys.filter(_ != storeProviderId).toSeq
      val providerIdsToUnload = reportActiveStoreInstance(storeProviderId, otherProviderIds)
-      providerIdsToUnload.foreach(unload(_))
+      providerIdsToUnload.foreach(id => {
+        loadedProviders.remove(id).foreach( provider => {


Already in the synchronized block on line 922.

micheal-o · 2025-04-24T17:22:19Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+          // Trigger maintenance thread to immediately do maintenance on and close the provider.
+          // Doing maintenance first allows us to do maintenance for a constantly-moving state
+          // store.
+          logInfo(log"Task thread trigger maintenance on " +


+1, that would help with investigation

micheal-o · 2025-04-24T18:12:14Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+          // Trigger maintenance thread to immediately do maintenance on and close the provider.
+          // Doing maintenance first allows us to do maintenance for a constantly-moving state
+          // store.
+          logInfo(log"Task thread trigger maintenance on " +


include trigger maintenance to close provider

micheal-o · 2025-04-24T18:17:09Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+   */
+  def unload(storeProviderId: StateStoreProviderId,
+             alreadyRemovedStoreFromLoadedProviders: Option[StateStoreProvider] = None): Unit = {
+    var toCloseProviders: List[StateStoreProvider] = Nil


why is this a list? we will only be closing one provider right?

micheal-o · 2025-04-24T18:32:55Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+      if (alreadyRemovedFromLoadedProviders) {
+        // If provider is already removed from loadedProviders, we MUST process
+        // this partition to close it, so we block until we can.
+        awaitProcessThisPartition(id)


lets make it clear that we are waiting for any possible ongoing maintenance on this partition.

also add comment about when this would have already been removed

micheal-o · 2025-04-24T18:32:59Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+        }
+      } catch {
+        case NonFatal(e) =>
+          logWarning(log"Error managing ${MDC(LogKeys.STATE_STORE_PROVIDER, provider)}, " +


nit: replace the word "managing"?

micheal-o · 2025-04-24T18:38:16Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+          // store.
+          logInfo(log"Task thread trigger maintenance on " +
+            log"provider=${MDC(LogKeys.STATE_STORE_PROVIDER, id)}")
+          doMaintenanceOnProvider(id, provider, alreadyRemovedFromLoadedProviders = true)


Lets rename this to submitMaintenanceWorkForProvider. To make it clear that this isn't blocking

liviazhu-db

Could you actually just while you're modifying this code also change all of the places where we have ${MDC(LogKeys.STATE_STORE_PROVIDER, provider)} (which logs the java object) to ${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)}? The first one doesn't give us much information, I think we really just want to see the ID in the logs

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

liviazhu-db · 2025-04-24T21:02:58Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+      val awaitingPartitionDuration = System.currentTimeMillis() - startTime
+      try {
+        provider.doMaintenance()
+        // If shouldRemoveFromLoadedProviders is false, we don't need to verify


Oops this comment needs to be updated

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

liviazhu-db

Also, not sure how possible it is, but could you try to write a test where the task submits maintenance work for the provider while it's already being processed by a maintenance thread (i.e. when we need to block in awaitProcessThisPartition)? If it's kind of hard to write a reliable unit test to test this, maybe a manual test would do if you hack around with the timing, just want to check that this functionality works as expected.

liviazhu-db

LGTM pending a nit

liviazhu-db · 2025-04-24T23:17:52Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

@@ -1009,14 +1013,46 @@ object StateStore extends Logging {

      val otherProviderIds = loadedProviders.keys.filter(_ != storeProviderId).toSeq
      val providerIdsToUnload = reportActiveStoreInstance(storeProviderId, otherProviderIds)
-      providerIdsToUnload.foreach(unload(_))
+      val taskContextIdLogLine = Option(TaskContext.get()).map { tc =>
+        log"taskId=${MDC(LogKeys.STATE_STORE_PROVIDER_ID, tc.taskAttemptId())}"


nit: I think LogKey is wrong here

anishshri-db · 2025-04-25T16:56:41Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+          // Trigger maintenance thread to immediately do maintenance on and close the provider.
+          // Doing maintenance first allows us to do maintenance for a constantly-moving state
+          // store.
+          logInfo(log"Task thread trigger maintenance to close " +


nit: Submitted maintenance from task thread

anishshri-db · 2025-04-25T16:59:27Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+        // with the coordinator as we know it definitely should be unloaded.
+        if (submittedFromTaskThread) {
+            unload(id, Some(provider))
+          } else if (!verifyIfStoreInstanceActive(id)) {


nit: indent seems off here ?

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

anishshri-db · 2025-04-28T20:58:53Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+  // Block until we can process this partition
+  private def awaitProcessThisPartition(
+      id: StateStoreProviderId,
+      storeConf: StateStoreConf): Boolean = {


Lets just pass the timeout value here ?

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

liviazhu-db

LGTM, thanks for making these changes!

liviazhu-db · 2025-04-29T17:34:35Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+          // (we've already removed it from loadedProviders)
+          true
+
+        case FromLoadedProviders =>


nit: comment?

+1 - lets comment this case too

anishshri-db · 2025-04-29T19:14:36Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+  private def awaitProcessThisPartition(
+      id: StateStoreProviderId,
+      timeoutMs: Long
+  ): Boolean = maintenanceThreadPoolLock synchronized  {


nit: can we move to line above ?

anishshri-db · 2025-04-29T19:16:35Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+      canProcessThisPartition = processThisPartition(id)
+      maintenanceThreadPoolLock.wait(timeoutMs)
+    }
+


can we add a log to indicate how much time we waited here ?

anishshri-db · 2025-04-29T19:49:54Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+    // Providers that couldn't be processed now and need to be added back to the queue
+    val providersToRequeue = new ArrayBuffer[(StateStoreProviderId, StateStoreProvider)]()
+
+    while (!unloadedProvidersToClose.isEmpty) {


Access to unloadedProvidersToClose is not truly thread safe here. Can we add a separate lock to guard access ?

unloadedProvidersToClose is a concurrent queue, so we don't need lock here. Since this is a while loop, if a new elem is added, we will see it. Even if we don't, not an issue, we will see it on next maintenance

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

ericm-db · 2025-04-30T05:24:36Z

cc @HeartSaVioR

micheal-o · 2025-05-02T04:41:55Z

common/utils/src/main/resources/error/error-conditions.json

@@ -4927,6 +4927,13 @@
    ],
    "sqlState" : "XXKST"
  },
+  "STATE_STORE_MAINTENANCE_TASK_TIMEOUT" : {


It seems we are currently not using this in this code. Couldn't find the usage.

micheal-o · 2025-05-02T04:43:06Z

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

+      .internal()
+      .doc("Timeout in seconds to wait for maintenance to process this partition.")
+      .timeConf(TimeUnit.SECONDS)
+      .createWithDefault(300L)


i think 5 mins is too large. Because this means a maintenance thread would be waiting for 5 mins for a partition. Instead of doing some other work. We can end up having many threads waiting, doing nothing for 5 mins. Lets set this lower cc @anishshri-db

micheal-o · 2025-05-02T04:43:21Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+    // Providers that couldn't be processed now and need to be added back to the queue
+    val providersToRequeue = new ArrayBuffer[(StateStoreProviderId, StateStoreProvider)]()
+
+    while (!unloadedProvidersToClose.isEmpty) {


unloadedProvidersToClose is a concurrent queue, so we don't need lock here. Since this is a while loop, if a new elem is added, we will see it. Even if we don't, not an issue, we will see it on next maintenance

micheal-o · 2025-05-02T04:43:25Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+    // Providers that couldn't be processed now and need to be added back to the queue
+    val providersToRequeue = new ArrayBuffer[(StateStoreProviderId, StateStoreProvider)]()
+
+    while (!unloadedProvidersToClose.isEmpty) {


nit: add comment about what providers will be in the unloadedProvidersToClose for future readers

micheal-o · 2025-05-02T04:43:37Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+      // Determine if we can process this partition based on the source
+      val canProcessThisPartition = source match {
+        case FromTaskThread =>
+          // Provider from task thread needs to wait for lock


nit: add additional comment that it needs to wait because there could already be an ongoing maintenance for that provider

micheal-o · 2025-05-02T04:44:18Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+          }
+
+          if (source != FromLoadedProviders) {
+            logInfo(log"Unloaded ${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)}")


we need this log too right if we unloaded for the FromLoadedProviders right?

micheal-o · 2025-05-02T04:44:22Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+          }
+        } catch {
+          case NonFatal(e) =>
+            logWarning(log"Error ${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)}, " +


nit: should we make this Error doing maintenance on or something?

micheal-o · 2025-05-02T04:44:30Z

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala

@@ -1718,6 +1845,240 @@ abstract class StateStoreSuiteBase[ProviderClass <: StateStoreProvider]
    assert(encoderSpec == deserializedEncoderSpec)
  }

+  test("SPARK-51596: task thread waits for ongoing maintenance to complete before" +


rename this, the task thread isn't the one waiting. Task thread just submits to maint threadpool

micheal-o · 2025-05-02T04:44:43Z

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala

+        }
+
+        // Now unblock the first maintenance operation
+        // This should allow the task's submitted maintenance to proceed


how do you know that the submited work actually waited for the ongoing maintenance and didn't happen concurrently? I don't see that verification here

micheal-o · 2025-05-02T04:44:51Z

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala

+
+        // Mark that task submitted maintenance
+        SignalingStateStoreProvider.taskSubmittedMaintenance = true
+


this test is actually not checking that the provider was queued. How are we sure it was queued?

micheal-o · 2025-05-02T04:50:39Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

+    var canProcessThisPartition = processThisPartition(id)
+    while (!canProcessThisPartition && System.currentTimeMillis() < endTime) {
+      canProcessThisPartition = processThisPartition(id)
+      maintenanceThreadPoolLock.wait(timeoutMs)


shouldn't the canProcessThisPartition be after the wait? Because with this current code if canProcessThisPartition becomes true, you will end up still waiting

liviazhu-db and others added 4 commits March 25, 2025 13:29

impl and unit test

248fa8e

fix unit tests

52085cb

remove store from loadedProviders from task thread

c294c9f

nits

3548ca3

github-actions bot added SQL STRUCTURED STREAMING labels Apr 15, 2025

ericm-db added 2 commits April 15, 2025 15:39

addressing shutdown

c8c0a4e

resolving conflicts

d42cc50

anishshri-db reviewed Apr 23, 2025

View reviewed changes

liviazhu-db reviewed Apr 24, 2025

View reviewed changes

micheal-o reviewed Apr 24, 2025

View reviewed changes

ericm-db added 3 commits April 24, 2025 13:52

addressing comments

642f3b4

imports

7f26f40

addressing Livia's comment

78af7fe

liviazhu-db reviewed Apr 24, 2025

View reviewed changes

ericm-db added 3 commits April 24, 2025 14:21

feedback

ca302f6

replacing log line

86f7e64

provider -> provider_id

c12a6ae

ericm-db requested a review from liviazhu-db April 24, 2025 21:27

ericm-db added 4 commits April 24, 2025 14:37

removing taskContext since it's null

1f7c3de

Merge remote-tracking branch 'upstream/master' into maintenance-changes

a74f935

removing taskContext since it's null

d17e476

adding back tc

c363483

liviazhu-db approved these changes Apr 24, 2025

View reviewed changes

ericm-db added 2 commits April 24, 2025 16:47

adding test

7079852

state_store_provider_id -> task_id

40fc8fb

anishshri-db reviewed Apr 25, 2025

View reviewed changes

liviazhu-db mentioned this pull request Apr 25, 2025

[SPARK-51596][SS] Fix concurrent StateStoreProvider maintenance and closing #50391

Closed

anishshri-db reviewed Apr 28, 2025

View reviewed changes

ericm-db added 2 commits April 28, 2025 16:00

queueing changes

e1ebd4a

seconds -> ms

db2506f

anishshri-db reviewed Apr 29, 2025

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala Outdated Show resolved Hide resolved

anishshri-db reviewed Apr 29, 2025

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala Outdated Show resolved Hide resolved

liviazhu-db suggested changes Apr 29, 2025

View reviewed changes

ericm-db added 5 commits April 29, 2025 09:39

boolean to enum and other feedback

7d5f168

thanks livia

dab9e8f

waiting for timeout

172ee83

wrapping in synch block

b051d5a

adding condition on submitMaintenance

7384a84

liviazhu-db approved these changes Apr 29, 2025

View reviewed changes

anishshri-db reviewed Apr 29, 2025

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala Outdated Show resolved Hide resolved

anishshri-db reviewed Apr 29, 2025

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala Outdated Show resolved Hide resolved

anishshri-db reviewed Apr 29, 2025

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala Outdated Show resolved Hide resolved

adding comment

46b300a

ericm-db force-pushed the maintenance-changes branch from 267c963 to 46b300a Compare April 29, 2025 21:05

anishshri-db approved these changes Apr 29, 2025

View reviewed changes

micheal-o reviewed May 2, 2025

View reviewed changes

ericm-db added 4 commits May 2, 2025 09:47

code feedback

776ad6b

adding parens

4946f9b

removing test

411cce2

adding test

9faf0a4

ericm-db force-pushed the maintenance-changes branch from 34220b6 to 9faf0a4 Compare May 2, 2025 20:57

adding log line

18a5d33


		// Mark that task submitted maintenance
		SignalingStateStoreProvider.taskSubmittedMaintenance = true

[SPARK-51596][SS] Fix concurrent StateStoreProvider maintenance and closing #50595

Are you sure you want to change the base?

[SPARK-51596][SS] Fix concurrent StateStoreProvider maintenance and closing #50595

Conversation

ericm-db commented Apr 15, 2025

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

liviazhu-db left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

liviazhu-db left a comment • edited Loading

Choose a reason for hiding this comment

liviazhu-db left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

liviazhu-db left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ericm-db commented Apr 30, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

liviazhu-db left a comment •

edited

Loading

liviazhu-db left a comment •

edited

Loading