[SPARK-51745] Enforce State Machine for RocksDBStateStore

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

@@ -5032,6 +5032,13 @@
     ],
     "sqlState" : "42802"
   },
+  "STATE_STORE_UPDATING_AFTER_TASK_COMPLETION" : {
+    "message" : [
+      "State store id=<stateStoreId> still in updating state after task completed. If using foreachBatch, ",
+      "verify that it consumes the entire dataframe and does not catch and suppress errors during dataframe iteration."
+    ],
+    "sqlState" : "XXKST"
+  },
   "STATE_STORE_VALUE_ROW_FORMAT_VALIDATION_FAILURE" : {
     "message" : [
       "The streaming query failed to validate written state for value row for stateStore=<stateStoreID>.",

common/utils/src/main/scala/org/apache/spark/internal/LogKey.scala

@@ -785,6 +785,7 @@ private[spark] object LogKeys {
   case object STATE_STORE_PROVIDER extends LogKey
   case object STATE_STORE_PROVIDER_ID extends LogKey
   case object STATE_STORE_PROVIDER_IDS extends LogKey
+  case object STATE_STORE_STATE extends LogKey
   case object STATE_STORE_VERSION extends LogKey
   case object STATS extends LogKey
   case object STATUS extends LogKey
@@ -867,6 +868,7 @@ private[spark] object LogKeys {
   case object TRAIN_VALIDATION_SPLIT_METRIC extends LogKey
   case object TRAIN_VALIDATION_SPLIT_METRICS extends LogKey
   case object TRANSFER_TYPE extends LogKey
+  case object TRANSITION extends LogKey
   case object TREE_NODE extends LogKey
   case object TRIGGER_INTERVAL extends LogKey
   case object UI_ACLS extends LogKey

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

@@ -2260,6 +2260,13 @@ object SQLConf {
       .timeConf(TimeUnit.SECONDS)
       .createWithDefault(300L)
 
+  val STATE_STORE_MAINTENANCE_PROCESSING_TIMEOUT =
+    buildConf("spark.sql.streaming.stateStore.maintenanceProcessingTimeout")
+      .internal()
+      .doc("Timeout in seconds to wait for maintenance to process this partition.")
+      .timeConf(TimeUnit.SECONDS)
+      .createWithDefault(30L)
+
   val STATE_SCHEMA_CHECK_ENABLED =
     buildConf("spark.sql.streaming.stateStore.stateSchemaCheck")
       .doc("When true, Spark will validate the state schema against schema on existing state and " +
@@ -6018,6 +6025,9 @@ class SQLConf extends Serializable with Logging with SqlApiConf {
 
   def stateStoreMaintenanceShutdownTimeout: Long = getConf(STATE_STORE_MAINTENANCE_SHUTDOWN_TIMEOUT)
 
+  def stateStoreMaintenanceProcessingTimeout: Long =
+    getConf(STATE_STORE_MAINTENANCE_PROCESSING_TIMEOUT)
+
   def stateStoreMinDeltasForSnapshot: Int = getConf(STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT)
 
   def stateStoreFormatValidationEnabled: Boolean = getConf(STATE_STORE_FORMAT_VALIDATION_ENABLED)

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala

@@ -126,6 +126,7 @@ abstract class StatePartitionReaderBase(
         stateStoreColFamilySchema.keyStateEncoderSpec.get,
         useMultipleValuesPerKey = useMultipleValuesPerKey,
         isInternal = isInternal)
+      store.abort()
     }
     provider
   }

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

@@ -31,7 +31,7 @@ import org.apache.commons.io.IOUtils
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs._
 
-import org.apache.spark.{SparkConf, SparkEnv, SparkException}
+import org.apache.spark.{SparkConf, SparkEnv, SparkException, TaskContext}
 import org.apache.spark.internal.{Logging, LogKeys, MDC, MessageWithContext}
 import org.apache.spark.io.CompressionCodec
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
@@ -89,6 +89,8 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
 
     override def abort(): Unit = {}
 
+    override def release(): Unit = {}
+
     override def toString(): String = {
       s"HDFSReadStateStore[id=(op=${id.operatorId},part=${id.partitionId}),dir=$baseDir]"
     }
@@ -112,6 +114,19 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
     case object UPDATING extends STATE
     case object COMMITTED extends STATE
     case object ABORTED extends STATE
+    case object RELEASED extends STATE
+
+    Option(TaskContext.get()).foreach { ctxt =>
+      ctxt.addTaskCompletionListener[Unit]( ctx => {
+        if (state == UPDATING) {
+          abort()
+          // Only throw error if the task is not already failed or interrupted
+          if (!ctx.isFailed() && !ctx.isInterrupted()) {
+            throw StateStoreErrors.stateStoreUpdatingAfterTaskCompletion(id)
+          }
+        }
+      })
+    }
 
     private val newVersion = version + 1
     @volatile private var state: STATE = UPDATING
@@ -194,6 +209,10 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
         log"for ${MDC(LogKeys.STATE_STORE_PROVIDER, this)}")
     }
 
+    override def release(): Unit = {
+      state = RELEASED
+    }
+
     /**
      * Get an iterator of all the store data.
      * This can be called only after committing all the updates made in the current thread.
@@ -953,7 +972,8 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
    * @param endVersion   checkpoint version to end with
    * @return [[HDFSBackedStateStore]]
    */
-  override def replayStateFromSnapshot(snapshotVersion: Long, endVersion: Long): StateStore = {
+  override def replayStateFromSnapshot(
+      snapshotVersion: Long, endVersion: Long, readOnly: Boolean): StateStore = {
     val newMap = replayLoadedMapFromSnapshot(snapshotVersion, endVersion)
     logInfo(log"Retrieved snapshot at version " +
       log"${MDC(LogKeys.STATE_STORE_VERSION, snapshotVersion)} and apply delta files to version " +