[SPARK-54063][SS] Trigger snapshot for next batch when upload lag

## What changes were proposed in this pull request?
We are adding a functionality for StateStoreProvider to force creating a snapshot on commit when creating a StateStore when it detects that the corresponding StateStoreProvider is lagging in creating snapshot (i.e. having too many changelog). This applies to both RocksDB and HDFS. This feature is only enabled when `STATE_STORE_COORDINATOR_REPORT_SNAPSHOT_UPLOAD_LAG` is true as well (introduced in [PR 50123](https://github.com/apache/spark/pull/50123/files#diff-7c577967a171f51523afddd7a5eca49806432fae2da8dc35c114d9699e6e3e40))

We also added a StateStoreCustomMetric `CUSTOM_METRIC_FORCE_SNAPSHOT` to track the number of forced snapshot

The changes made in this PR includes
-  `StateStoreCoordinator` now keeps track of a list of lagging state stores and update it accordingly on receiveing `LogLaggingStateStores` (update the set) and `DeactivateInstances` (remove all lagging stores from set)
- Change `ReportActiveInstance` in StateStoreCoordinator to also return `shouldForceSnapshot`, which indicates if current `StateStore` should create a snapshot on next commit
- Add a new param `forceSnapshotOnCommit: Boolean  =  false` to StateStoreProvider function `getStore`, and its inherited provider can create state stores that will force a snapshot on commit
- When StateStore singleton `get()` is called, it will call `ReportActiveInstance` to check if it should create a snapshot in the next commit.

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

No

### How was this patch tested?

Added integration test in StateStoreCoordinatorSuite, where it tests if RocksDBStateStore and HDFSBackedStateStore can both auto-trigger a snapshot on commit when delay in snapshot upload is detected

Also added unit test in StateStoreBaseSuite to test if the metric is populated correctly when a snapshot is forced.

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

Yes

Closes #52773 from zifeif2/trigger-snapshot.

Authored-by: Ubuntu <zifei.feng@your.hostname.com>
Signed-off-by: Anish Shrigondekar <anish.shrigondekar@databricks.com>

Author: Ubuntu

common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java

@@ -787,6 +787,7 @@ public enum LogKeys implements LogKey {
   STREAM_CHUNK_ID,
   STREAM_ID,
   STREAM_NAME,
+  STREAM_SHOULD_FORCE_SNAPSHOT,
   SUBMISSION_ID,
   SUBSAMPLING_RATE,
   SUB_QUERY,

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

@@ -2710,6 +2710,18 @@ object SQLConf {
       .checkValue(k => k >= 0, "Must be greater than or equal to 0")
       .createWithDefault(if (Utils.isTesting) 1 else 0)
 
+  val STATE_STORE_FORCE_SNAPSHOT_UPLOAD_ON_LAG =
+    buildConf("spark.sql.streaming.stateStore.forceSnapshotUploadOnLag")
+      .internal()
+      .doc(
+        "When enabled, state stores with lagging snapshot uploads will automatically trigger " +
+        "a snapshot on the next commit. Requires spark.sql.streaming.stateStore.coordinator" +
+        "ReportSnapshotUploadLag to be true."
+      )
+      .version("4.2.0")
+      .booleanConf
+      .createWithDefault(false)
+
   val STATEFUL_SHUFFLE_PARTITIONS_INTERNAL =
     buildConf("spark.sql.streaming.internal.stateStore.partitions")
       .doc("WARN: This config is used internally and is not intended to be user-facing. This " +
@@ -6921,6 +6933,17 @@ class SQLConf extends Serializable with Logging with SqlApiConf {
   def stateStoreRowChecksumReadVerificationRatio: Long =
     getConf(STATE_STORE_ROW_CHECKSUM_READ_VERIFICATION_RATIO)
 
+  def stateStoreForceSnapshotUploadOnLag: Boolean = {
+    val value = getConf(STATE_STORE_FORCE_SNAPSHOT_UPLOAD_ON_LAG)
+    if (value && !stateStoreCoordinatorReportSnapshotUploadLag) {
+      throw new IllegalArgumentException(
+        "spark.sql.streaming.stateStore.forceSnapshotUploadOnLag can only be true if " +
+        "spark.sql.streaming.stateStore.coordinatorReportSnapshotUploadLag is also true."
+      )
+    }
+    value
+  }
+
   def checkpointLocation: Option[String] = getConf(CHECKPOINT_LOCATION)
 
   def checkpointFileChecksumEnabled: Boolean = getConf(STREAMING_CHECKPOINT_FILE_CHECKSUM_ENABLED)

sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/streaming/ClientStreamingQuerySuite.scala

@@ -101,6 +101,7 @@ class ClientStreamingQuerySuite extends QueryTest with RemoteSparkSession with L
         assert(lastProgress.stateOperators.nonEmpty)
         assert(
           lastProgress.stateOperators.head.customMetrics.keySet().asScala == Set(
+            "forceSnapshotCount",
             "loadedMapCacheHitCount",
             "loadedMapCacheMissCount",
             "numSnapshotsAutoRepaired",

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

@@ -109,7 +109,8 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
   /** Implementation of [[StateStore]] API which is backed by an HDFS-compatible file system */
   class HDFSBackedStateStore(
       val version: Long,
-      private val mapToUpdate: HDFSBackedStateStoreMap)
+      private val mapToUpdate: HDFSBackedStateStoreMap,
+      shouldForceSnapshot: Boolean = false)
     extends StateStore {
 
     /** Trait and classes representing the internal state of the store */
@@ -196,7 +197,7 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
     override def commit(): Long = {
       try {
         verify(state == UPDATING, "Cannot commit after already committed or aborted")
-        commitUpdates(newVersion, mapToUpdate, compressedStream)
+        commitUpdates(newVersion, mapToUpdate, compressedStream, shouldForceSnapshot)
         state = COMMITTED
         logInfo(log"Committed version ${MDC(LogKeys.COMMITTED_VERSION, newVersion)} " +
           log"for ${MDC(LogKeys.STATE_STORE_PROVIDER, this)} to file " +
@@ -254,7 +255,8 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
       val customMetrics = metricsFromProvider.flatMap { case (name, value) =>
         // just allow searching from list cause the list is small enough
         supportedCustomMetrics.find(_.name == name).map(_ -> value)
-      } + (metricStateOnCurrentVersionSizeBytes -> SizeEstimator.estimate(mapToUpdate))
+      } + (metricStateOnCurrentVersionSizeBytes -> SizeEstimator.estimate(mapToUpdate)) +
+        (metricForceSnapshot -> (if (shouldForceSnapshot) 1L else 0L))
 
       val instanceMetrics = Map(
         instanceMetricSnapshotLastUpload.withNewId(
@@ -317,16 +319,21 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
   }
 
   /** Get the state store for making updates to create a new `version` of the store. */
-  override def getStore(version: Long, uniqueId: Option[String] = None): StateStore = {
+  override def getStore(
+      version: Long,
+      uniqueId: Option[String] = None,
+      forceSnapshotOnCommit: Boolean = false): StateStore = {
     if (uniqueId.isDefined) {
       throw StateStoreErrors.stateStoreCheckpointIdsNotSupported(
         "HDFSBackedStateStoreProvider does not support checkpointFormatVersion > 1 " +
         "but a state store checkpointID is passed in")
     }
     val newMap = getLoadedMapForStore(version)
     logInfo(log"Retrieved version ${MDC(LogKeys.STATE_STORE_VERSION, version)} " +
-      log"of ${MDC(LogKeys.STATE_STORE_PROVIDER, HDFSBackedStateStoreProvider.this)} for update")
-    new HDFSBackedStateStore(version, newMap)
+      log"of ${MDC(LogKeys.STATE_STORE_PROVIDER, HDFSBackedStateStoreProvider.this)} " +
+      log"for update, forceSnapshotOnCommit=" +
+      log"${MDC(LogKeys.STREAM_SHOULD_FORCE_SNAPSHOT, forceSnapshotOnCommit)}")
+    new HDFSBackedStateStore(version, newMap, forceSnapshotOnCommit)
   }
 
   /** Get the state store for reading to specific `version` of the store. */
@@ -446,8 +453,7 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
 
   override def supportedCustomMetrics: Seq[StateStoreCustomMetric] = {
     metricStateOnCurrentVersionSizeBytes :: metricLoadedMapCacheHit :: metricLoadedMapCacheMiss ::
-    metricNumSnapshotsAutoRepaired ::
-      Nil
+    metricNumSnapshotsAutoRepaired :: metricForceSnapshot :: Nil
   }
 
   override def supportedInstanceMetrics: Seq[StateStoreInstanceMetric] =
@@ -524,6 +530,10 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
     StateStoreCustomSumMetric("numSnapshotsAutoRepaired",
     "number of snapshots that were automatically repaired during store load")
 
+  private lazy val metricForceSnapshot: StateStoreCustomMetric =
+    StateStoreCustomSumMetric("forceSnapshotCount",
+    "number of stores that had forced snapshot")
+
   private lazy val instanceMetricSnapshotLastUpload: StateStoreInstanceMetric =
     StateStoreSnapshotLastUploadInstanceMetric()
 
@@ -532,9 +542,13 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
   private def commitUpdates(
       newVersion: Long,
       map: HDFSBackedStateStoreMap,
-      output: DataOutputStream): Unit = {
+      output: DataOutputStream,
+      shouldForceSnapshot: Boolean = false): Unit = {
     synchronized {
       finalizeDeltaFile(output)
+      if (shouldForceSnapshot) {
+        writeSnapshotFile(newVersion, map, "commit")
+      }
       putStateIntoStateCacheMap(newVersion, map)
     }
   }

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

@@ -1469,13 +1469,17 @@ class RocksDB(
    * - Create a RocksDB checkpoint in a new local dir
    * - Sync the checkpoint dir files to DFS
    */
-  def commit(): (Long, StateStoreCheckpointInfo) = {
+  def commit(forceSnapshot: Boolean = false): (Long, StateStoreCheckpointInfo) = {
     commitLatencyMs.clear()
     updateMemoryUsageIfNeeded()
     val newVersion = loadedVersion + 1
     try {
       logInfo(log"Flushing updates for ${MDC(LogKeys.VERSION_NUM, newVersion)}")
 
+      if (forceSnapshot) {
+        shouldForceSnapshot.set(true)
+      }
+
       var snapshot: Option[RocksDBSnapshot] = None
       if (shouldCreateSnapshot() || shouldForceSnapshot.get()) {
         val (newSnapshot, snapshotLatency) = createSnapshot(newVersion, sessionStateStoreCkptId)

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

@@ -46,7 +46,8 @@ private[sql] class RocksDBStateStoreProvider
   class RocksDBStateStore(
       lastVersion: Long,
       private[RocksDBStateStoreProvider] val stamp: Long,
-      private[RocksDBStateStoreProvider] var readOnly: Boolean) extends StateStore {
+      private[RocksDBStateStoreProvider] var readOnly: Boolean,
+      private[RocksDBStateStoreProvider] var forceSnapshotOnCommit: Boolean) extends StateStore {
 
     private sealed trait OPERATION
     private case object UPDATE extends OPERATION
@@ -448,7 +449,7 @@ private[sql] class RocksDBStateStoreProvider
       validateState(UPDATING)
       try {
         stateMachine.verifyStamp(stamp)
-        val (newVersion, newCheckpointInfo) = rocksDB.commit()
+        val (newVersion, newCheckpointInfo) = rocksDB.commit(forceSnapshotOnCommit)
         checkpointInfo = Some(newCheckpointInfo)
         storedMetrics = rocksDB.metricsOpt
         validateAndTransitionState(COMMIT)
@@ -568,7 +569,8 @@ private[sql] class RocksDBStateStoreProvider
           CUSTOM_METRIC_NUM_INTERNAL_COL_FAMILIES_KEYS -> rocksDBMetrics.numInternalKeys,
           CUSTOM_METRIC_NUM_EXTERNAL_COL_FAMILIES -> internalColFamilyCnt(),
           CUSTOM_METRIC_NUM_INTERNAL_COL_FAMILIES -> externalColFamilyCnt(),
-          CUSTOM_METRIC_NUM_SNAPSHOTS_AUTO_REPAIRED -> rocksDBMetrics.numSnapshotsAutoRepaired
+          CUSTOM_METRIC_NUM_SNAPSHOTS_AUTO_REPAIRED -> rocksDBMetrics.numSnapshotsAutoRepaired,
+          CUSTOM_METRIC_FORCE_SNAPSHOT -> (if (forceSnapshotOnCommit) 1L else 0L)
         ) ++ rocksDBMetrics.zipFileBytesUncompressed.map(bytes =>
           Map(CUSTOM_METRIC_ZIP_FILE_BYTES_UNCOMPRESSED -> bytes)).getOrElse(Map())
 
@@ -716,13 +718,15 @@ private[sql] class RocksDBStateStoreProvider
    * @param uniqueId Optional unique identifier for checkpoint
    * @param readOnly Whether to open the store in read-only mode
    * @param existingStore Optional existing store to reuse instead of creating a new one
+   * @param forceSnapshotOnCommit Whether to force a snapshot upload on commit
    * @return The loaded state store
    */
   private def loadStateStore(
       version: Long,
       uniqueId: Option[String] = None,
       readOnly: Boolean,
-      existingStore: Option[RocksDBStateStore] = None): StateStore = {
+      existingStore: Option[RocksDBStateStore] = None,
+      forceSnapshotOnCommit: Boolean = false): StateStore = {
     var acquiredStamp: Option[Long] = None
     var storeLoaded = false
     try {
@@ -761,11 +765,12 @@ private[sql] class RocksDBStateStoreProvider
         case Some(store: RocksDBStateStore) =>
           // Mark store as being used for write operations
           store.readOnly = readOnly
+          store.forceSnapshotOnCommit = forceSnapshotOnCommit
           store
         case None =>
           // Create new store instance. The stamp should be defined
           // in this case
-          new RocksDBStateStore(version, stamp.get, readOnly)
+          new RocksDBStateStore(version, stamp.get, readOnly, forceSnapshotOnCommit)
       }
       storeLoaded = true
       store
@@ -792,23 +797,34 @@ private[sql] class RocksDBStateStoreProvider
   }
 
   override def getStore(
-      version: Long, uniqueId: Option[String] = None): StateStore = {
-    loadStateStore(version, uniqueId, readOnly = false)
+      version: Long,
+      uniqueId: Option[String] = None,
+      forceSnapshotOnCommit: Boolean = false): StateStore = {
+    loadStateStore(
+      version,
+      uniqueId,
+      readOnly = false,
+      forceSnapshotOnCommit = forceSnapshotOnCommit
+    )
   }
 
   override def upgradeReadStoreToWriteStore(
       readStore: ReadStateStore,
       version: Long,
-      uniqueId: Option[String] = None): StateStore = {
+      uniqueId: Option[String] = None,
+      forceSnapshotOnCommit: Boolean = false): StateStore = {
     assert(version == readStore.version,
       s"Can only upgrade readStore to writeStore with the same version," +
         s" readStoreVersion: ${readStore.version}, writeStoreVersion: ${version}")
     assert(this.stateStoreId == readStore.id, "Can only upgrade readStore to writeStore with" +
       " the same stateStoreId")
     assert(readStore.isInstanceOf[RocksDBStateStore], "Can only upgrade state store if it is a " +
       "RocksDBStateStore")
-    loadStateStore(version, uniqueId, readOnly = false, existingStore =
-      Some(readStore.asInstanceOf[RocksDBStateStore]))
+    loadStateStore(version,
+      uniqueId,
+      readOnly = false,
+      existingStore = Some(readStore.asInstanceOf[RocksDBStateStore]),
+      forceSnapshotOnCommit = forceSnapshotOnCommit)
   }
 
   override def getReadStore(
@@ -932,7 +948,7 @@ private[sql] class RocksDBStateStoreProvider
           endVersion,
           snapshotVersionStateStoreCkptId,
           endVersionStateStoreCkptId)
-        new RocksDBStateStore(endVersion, stamp, readOnly)
+        new RocksDBStateStore(endVersion, stamp, readOnly, forceSnapshotOnCommit = false)
       } catch {
         case e: Throwable =>
           stateMachine.releaseStamp(stamp)
@@ -1283,6 +1299,9 @@ object RocksDBStateStoreProvider {
     "rocksdbNumSnapshotsAutoRepaired",
     "RocksDB: number of snapshots that were automatically repaired during store load")
 
+  val CUSTOM_METRIC_FORCE_SNAPSHOT = StateStoreCustomSumMetric(
+    "rocksdbForceSnapshotCount", "RocksDB: number of stores that had forced snapshot on commit")
+
   val ALL_CUSTOM_METRICS = Seq(
     CUSTOM_METRIC_SST_FILE_SIZE, CUSTOM_METRIC_GET_TIME, CUSTOM_METRIC_PUT_TIME,
     CUSTOM_METRIC_FLUSH_TIME, CUSTOM_METRIC_COMMIT_COMPACT_TIME,
@@ -1297,7 +1316,8 @@ object RocksDBStateStoreProvider {
     CUSTOM_METRIC_PINNED_BLOCKS_MEM_USAGE, CUSTOM_METRIC_NUM_INTERNAL_COL_FAMILIES_KEYS,
     CUSTOM_METRIC_NUM_EXTERNAL_COL_FAMILIES, CUSTOM_METRIC_NUM_INTERNAL_COL_FAMILIES,
     CUSTOM_METRIC_LOAD_FROM_SNAPSHOT_TIME, CUSTOM_METRIC_LOAD_TIME, CUSTOM_METRIC_REPLAY_CHANGE_LOG,
-    CUSTOM_METRIC_NUM_REPLAY_CHANGE_LOG_FILES, CUSTOM_METRIC_NUM_SNAPSHOTS_AUTO_REPAIRED)
+    CUSTOM_METRIC_NUM_REPLAY_CHANGE_LOG_FILES, CUSTOM_METRIC_NUM_SNAPSHOTS_AUTO_REPAIRED,
+    CUSTOM_METRIC_FORCE_SNAPSHOT)
 
   val CUSTOM_INSTANCE_METRIC_SNAPSHOT_LAST_UPLOADED = StateStoreSnapshotLastUploadInstanceMetric()