[SPARK-54078][SS] New test for StateStoreSuite SPARK-40492: maintenance before unload and remove infra from old test

### What changes were proposed in this pull request?

The `SPARK-40492: maintenance before unload` test is occasionally flaky even after refactors in #52783. This PR rewrites the test by using a new state store provider that counts number of times maintenance is triggered to test the functionality more directly. It also removes the hooks that were only used by the old test.

### Why are the changes needed?

Deflake the test

### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

Unit test, ran over 60 times

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

Generated-by: Claude 4

Closes #52877 from liviazhu/statestore-flakiness-2.

Authored-by: Livia Zhu <[email protected]>
Signed-off-by: Anish Shrigondekar <[email protected]>

Author: liviazhu

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

@@ -1424,12 +1424,6 @@ object StateStore extends Logging {
     }
   }
 
-  // Pause maintenance for testing purposes only.
-  @volatile private var maintenancePaused: Boolean = false
-  private[spark] def setMaintenancePaused(maintPaused: Boolean): Unit = {
-      maintenancePaused = maintPaused
-  }
-
   /**
    * Execute background maintenance task in all the loaded store providers if they are still
    * the active instances according to the coordinator.
@@ -1439,10 +1433,6 @@ object StateStore extends Logging {
     if (SparkEnv.get == null) {
       throw new IllegalStateException("SparkEnv not active, cannot do maintenance on StateStores")
     }
-    if (maintenancePaused) {
-      logDebug("Maintenance paused")
-      return
-    }
 
     // Providers that couldn't be processed now and need to be added back to the queue
     val providersToRequeue = new ArrayBuffer[(StateStoreProviderId, StateStoreProvider)]()

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

@@ -254,6 +254,28 @@ private object FakeStateStoreProviderWithMaintenanceError {
   val errorOnMaintenance = new AtomicBoolean(false)
 }
 
+/**
+ * A fake StateStoreProvider for testing maintenance before unload.
+ * Extends HDFSBackedStateStoreProvider to get actual store functionality,
+ * but tracks the number of times doMaintenance is called.
+ */
+class MaintenanceCountingStateStoreProvider extends HDFSBackedStateStoreProvider {
+  import MaintenanceCountingStateStoreProvider._
+
+  override def doMaintenance(): Unit = {
+    maintenanceCallCount.incrementAndGet()
+    super.doMaintenance()
+  }
+}
+
+private object MaintenanceCountingStateStoreProvider {
+  val maintenanceCallCount = new java.util.concurrent.atomic.AtomicInteger(0)
+
+  def reset(): Unit = {
+    maintenanceCallCount.set(0)
+  }
+}
+
 @ExtendedSQLTest
 class StateStoreSuite extends StateStoreSuiteBase[HDFSBackedStateStoreProvider]
   with SharedSparkSession
@@ -1013,77 +1035,65 @@ class StateStoreSuite extends StateStoreSuiteBase[HDFSBackedStateStoreProvider]
 
   // Ensure that maintenance is called before unloading
   test("SPARK-40492: maintenance before unload") {
+    // Reset the maintenance call counter
+    MaintenanceCountingStateStoreProvider.reset()
+
     val conf = new SparkConf()
       .setMaster("local")
       .setAppName("SPARK-40492")
     val opId = 0
     val dir1 = newDir()
     val storeProviderId1 = StateStoreProviderId(StateStoreId(dir1, opId, 0), UUID.randomUUID)
-    val sqlConf = getDefaultSQLConf(SQLConf.STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT.defaultValue.get,
-      SQLConf.MAX_BATCHES_TO_RETAIN_IN_MEMORY.defaultValue.get)
-    sqlConf.setConf(SQLConf.MIN_BATCHES_TO_RETAIN, 10)
-    sqlConf.setConf(SQLConf.STREAMING_MAINTENANCE_INTERVAL, 10L)
+    val sqlConf = getDefaultSQLConf(
+      SQLConf.STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT.defaultValue.get,
+      SQLConf.MAX_BATCHES_TO_RETAIN_IN_MEMORY.defaultValue.get
+    )
+    sqlConf.setConf(SQLConf.STREAMING_MAINTENANCE_INTERVAL, 5.seconds.toMillis)
+    sqlConf.setConf(SQLConf.STATE_STORE_PROVIDER_CLASS,
+      classOf[MaintenanceCountingStateStoreProvider].getName)
     val storeConf = StateStoreConf(sqlConf)
     val hadoopConf = new Configuration()
 
-    var latestStoreVersion = 0
-
-    def generateStoreVersions(): Unit = {
-      for (i <- 1 to 20) {
-        val store = StateStore.get(storeProviderId1, keySchema, valueSchema,
-          NoPrefixKeyStateEncoderSpec(keySchema),
-          latestStoreVersion, None, None, useColumnFamilies = false, storeConf, hadoopConf)
-        put(store, "a", 0, i)
-        store.commit()
-        latestStoreVersion += 1
-      }
-    }
-
     val timeoutDuration = 1.minute
 
     quietly {
       withSpark(SparkContext.getOrCreate(conf)) { sc =>
         withCoordinatorRef(sc) { coordinatorRef =>
           require(!StateStore.isMaintenanceRunning, "StateStore is unexpectedly running")
 
-          // Generate sufficient versions of store for snapshots
-          generateStoreVersions()
+          // Load the store
+          StateStore.get(storeProviderId1, keySchema, valueSchema,
+            NoPrefixKeyStateEncoderSpec(keySchema),
+            0, None, None, useColumnFamilies = false, storeConf, hadoopConf)
+
+          // Ensure the store is loaded
           eventually(timeout(timeoutDuration)) {
-            // Store should have been reported to the coordinator
             assert(coordinatorRef.getLocation(storeProviderId1).nonEmpty,
               "active instance was not reported")
-            // Background maintenance should clean up and generate snapshots
-            assert(StateStore.isMaintenanceRunning, "Maintenance task is not running")
-            // Some snapshots should have been generated
-            tryWithProviderResource(newStoreProvider(storeProviderId1.storeId)) { provider =>
-              val snapshotVersions = (1 to latestStoreVersion).filter { version =>
-                fileExists(provider, version, isSnapshot = true)
-              }
-              assert(snapshotVersions.nonEmpty, "no snapshot file found")
-            }
+            assert(StateStore.isLoaded(storeProviderId1), "Store is not loaded")
           }
-          // Pause maintenance
-          StateStore.setMaintenancePaused(true)
 
-          // Generate more versions such that there is another snapshot.
-          generateStoreVersions()
+          // Record the current maintenance call count before deactivation
+          val maintenanceCountBeforeDeactivate =
+            MaintenanceCountingStateStoreProvider.maintenanceCallCount.get()
 
-          // If driver decides to deactivate all stores related to a query run,
-          // then this instance should be unloaded.
+          // Deactivate the store instance - this should trigger maintenance before unload
           coordinatorRef.deactivateInstances(storeProviderId1.queryRunId)
 
-          // Resume maintenance which should unload the deactivated store
-          StateStore.setMaintenancePaused(false)
+          // Wait for the store to be unloaded
           eventually(timeout(timeoutDuration)) {
-            assert(!StateStore.isLoaded(storeProviderId1))
+            assert(!StateStore.isLoaded(storeProviderId1), "Store was not unloaded")
           }
 
-          // Ensure the earliest delta file should be cleaned up during unload.
-          tryWithProviderResource(newStoreProvider(storeProviderId1.storeId)) { provider =>
-            eventually(timeout(timeoutDuration)) {
-              assert(!fileExists(provider, 1, isSnapshot = false), "earliest file not deleted")
-            }
-          }
+          // Get the maintenance count after unload
+          val maintenanceCountAfterUnload =
+            MaintenanceCountingStateStoreProvider.maintenanceCallCount.get()
+
+          // Ensure that maintenance was called at least one more time during unload
+          assert(maintenanceCountAfterUnload > maintenanceCountBeforeDeactivate,
+            s"Maintenance should be called before unload. " +
+              s"Before: $maintenanceCountBeforeDeactivate, " +
+              s"After: $maintenanceCountAfterUnload")
         }
       }
     }