[SPARK-51252][SS] Add instance metrics for last uploaded snapshot version in HDFS State Stores

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

SPARK-51252

Similar to SPARK-51097, This PR sets up instance-specific metrics (`SnapshotLastUploaded.partition_<partition id>_<state store name>` to be precise) in the executor side and publishes them through StreamingQueryProgress. The main difference is that this PR adds the metrics to HDFS state stores to have similar observability to RocksDB state stores.

There are two non-test differences compared to the RocksDB PR:

- Add tracking of lastUploadedSnapshotVersion to HDFS state stores, which is updated in the state store object instead of RocksDB. Also updates at the end of writeUpdates instead of uploadSnapshot.
- Metric creation is done in the state store provider class instead of object

The fixes for SparkUI does not affect this PR, since that was an issue with the underlying instance metrics collection process. I still re-verified that these metrics do not show up on SparkUI (screenshot below)

As for the test changes:

- Modified the RocksDB tests slightly to process some additional data in the queries, because the version difference check used in HDFS and RocksDB state stores uses > and >= respectively. The general structure of the tests remain the same.

### Why are the changes needed?

We would like to introduce a similar level of observability into other types of state stores. As for the metrics themselves, it would be helpful in identifying performance degradation issues behind maintenance tasks and more as described in SPARK-51097.

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

There will be some new metrics displayed from StreamingQueryProgress:
```
Streaming query made progress: {
  ...
  "stateOperators" : [ {
    ...
    "customMetrics" : {
      ...
      "SnapshotLastUploaded.partition_0_default" : 2,
      "SnapshotLastUploaded.partition_12_default" : 10,
      "SnapshotLastUploaded.partition_8_default" : 10,
      ...
    }
  } ],
  "sources" : ...,
  "sink" : ...
}
```
Similar to its RocksDB variant, the amount of metrics reported is also tuned using `STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT`, and will by default report 5 instance metrics.

### How was this patch tested?

Four new tests are added into a new suite StateStoreInstanceMetricSuite, making eight tests in total for both state stores. Future instance metrics can be included into this suite for testing and verification.

The first two tests execute a dedup streaming query and verifies metrics are properly filtered + updated through the StreamingQueryProgress logs, but with different StateStore providers that skip maintenance tasks for specific partitions.

The other two tests execute a join streaming query, which contains four state stores per partition instead of one. These two tests verifies metrics are properly collected and filtered as well.

Verified the metrics do not show up on Spark UI:
![image](https://github.com/user-attachments/assets/c3c722fe-3bca-44a3-8f8d-85ed865def05)

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

No.

Closes #50030 from zecookiez/SPARK-51252-hdfs.

Lead-authored-by: Zeyu Chen <zycm03@gmail.com>
Co-authored-by: Zeyu Chen <ZYCM03@gmail.com>
Signed-off-by: Jungtaek Lim <kabhwan.opensource@gmail.com>

Author: zecookiez

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

@@ -103,7 +103,8 @@ class ClientStreamingQuerySuite extends QueryTest with RemoteSparkSession with L
           lastProgress.stateOperators.head.customMetrics.keySet().asScala == Set(
             "loadedMapCacheHitCount",
             "loadedMapCacheMissCount",
-            "stateOnCurrentVersionSizeBytes"))
+            "stateOnCurrentVersionSizeBytes",
+            "SnapshotLastUploaded.partition_0_default"))
         assert(lastProgress.sources.nonEmpty)
         assert(lastProgress.sink.description == "MemorySink")
         assert(lastProgress.observedMetrics.isEmpty)

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

@@ -20,7 +20,7 @@ package org.apache.spark.sql.execution.streaming.state
 import java.io._
 import java.util
 import java.util.Locale
-import java.util.concurrent.atomic.LongAdder
+import java.util.concurrent.atomic.{AtomicLong, LongAdder}
 
 import scala.collection.mutable
 import scala.jdk.CollectionConverters._
@@ -219,7 +219,17 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
         supportedCustomMetrics.find(_.name == name).map(_ -> value)
       } + (metricStateOnCurrentVersionSizeBytes -> SizeEstimator.estimate(mapToUpdate))
 
-      StateStoreMetrics(mapToUpdate.size(), metricsFromProvider("memoryUsedBytes"), customMetrics)
+      val instanceMetrics = Map(
+        instanceMetricSnapshotLastUpload.withNewId(
+          stateStoreId.partitionId, stateStoreId.storeName) -> lastUploadedSnapshotVersion.get()
+      )
+
+      StateStoreMetrics(
+        mapToUpdate.size(),
+        metricsFromProvider("memoryUsedBytes"),
+        customMetrics,
+        instanceMetrics
+      )
     }
 
     override def getStateStoreCheckpointInfo(): StateStoreCheckpointInfo = {
@@ -386,6 +396,9 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
       Nil
   }
 
+  override def supportedInstanceMetrics: Seq[StateStoreInstanceMetric] =
+    Seq(instanceMetricSnapshotLastUpload)
+
   private def toMessageWithContext: MessageWithContext = {
     log"HDFSStateStoreProvider[id = (op=${MDC(LogKeys.OP_ID, stateStoreId.operatorId)}," +
       log"part=${MDC(LogKeys.PARTITION_ID, stateStoreId.partitionId)})," +
@@ -419,6 +432,10 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
   private val loadedMapCacheHitCount: LongAdder = new LongAdder
   private val loadedMapCacheMissCount: LongAdder = new LongAdder
 
+  // This is updated when the maintenance task writes the snapshot file and read by the task
+  // thread. -1 represents no version has ever been uploaded.
+  private val lastUploadedSnapshotVersion: AtomicLong = new AtomicLong(-1L)
+
   private lazy val metricStateOnCurrentVersionSizeBytes: StateStoreCustomSizeMetric =
     StateStoreCustomSizeMetric("stateOnCurrentVersionSizeBytes",
       "estimated size of state only on current version")
@@ -431,6 +448,9 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
     StateStoreCustomSumMetric("loadedMapCacheMissCount",
       "count of cache miss on states cache in provider")
 
+  private lazy val instanceMetricSnapshotLastUpload: StateStoreInstanceMetric =
+    StateStoreSnapshotLastUploadInstanceMetric()
+
   private case class StoreFile(version: Long, path: Path, isSnapshot: Boolean)
 
   private def commitUpdates(
@@ -677,6 +697,8 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
     logInfo(log"Written snapshot file for version ${MDC(LogKeys.FILE_VERSION, version)} of " +
       log"${MDC(LogKeys.STATE_STORE_PROVIDER, this)} at ${MDC(LogKeys.FILE_NAME, targetFile)} " +
       log"for ${MDC(LogKeys.OP_TYPE, opType)}")
+    // Compare and update with the version that was just uploaded.
+    lastUploadedSnapshotVersion.updateAndGet(v => Math.max(version, v))
   }
 
   /**

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

@@ -19,36 +19,21 @@ package org.apache.spark.sql.execution.streaming.state
 
 import java.io.File
 
-import scala.concurrent.duration.DurationInt
-import scala.jdk.CollectionConverters.{MapHasAsScala, SetHasAsScala}
+import scala.jdk.CollectionConverters.SetHasAsScala
 
 import org.scalatest.time.{Minute, Span}
 
 import org.apache.spark.sql.execution.streaming.{MemoryStream, StreamingQueryWrapper}
-import org.apache.spark.sql.functions.{count, expr}
+import org.apache.spark.sql.functions.count
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.streaming._
 import org.apache.spark.sql.streaming.OutputMode.Update
 import org.apache.spark.util.Utils
 
-// SkipMaintenanceOnCertainPartitionsProvider is a test-only provider that skips running
-// maintenance for partitions 0 and 1 (these are arbitrary choices). This is used to test
-// snapshot upload lag can be observed through StreamingQueryProgress metrics.
-class SkipMaintenanceOnCertainPartitionsProvider extends RocksDBStateStoreProvider {
-  override def doMaintenance(): Unit = {
-    if (stateStoreId.partitionId == 0 || stateStoreId.partitionId == 1) {
-      return
-    }
-    super.doMaintenance()
-  }
-}
-
 class RocksDBStateStoreIntegrationSuite extends StreamTest
   with AlsoTestWithRocksDBFeatures {
   import testImplicits._
 
-  private val SNAPSHOT_LAG_METRIC_PREFIX = "SnapshotLastUploaded.partition_"
-
   testWithColumnFamilies("RocksDBStateStore",
     TestWithBothChangelogCheckpointingEnabledAndDisabled) { colFamiliesEnabled =>
     withTempDir { dir =>
@@ -286,277 +271,4 @@ class RocksDBStateStoreIntegrationSuite extends StreamTest
     assert(changelogVersionsPresent(dirForPartition0) == List(3L, 4L))
     assert(snapshotVersionsPresent(dirForPartition0).contains(5L))
   }
-
-  private def snapshotLagMetricName(
-      partitionId: Long,
-      storeName: String = StateStoreId.DEFAULT_STORE_NAME): String = {
-    s"$SNAPSHOT_LAG_METRIC_PREFIX${partitionId}_$storeName"
-  }
-
-  testWithChangelogCheckpointingEnabled(
-    "SPARK-51097: Verify snapshot lag metrics are updated correctly with RocksDBStateStoreProvider"
-  ) {
-    withSQLConf(
-      SQLConf.STATE_STORE_PROVIDER_CLASS.key -> classOf[RocksDBStateStoreProvider].getName,
-      SQLConf.STREAMING_MAINTENANCE_INTERVAL.key -> "100",
-      SQLConf.STREAMING_NO_DATA_PROGRESS_EVENT_INTERVAL.key -> "10",
-      SQLConf.STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT.key -> "1",
-      SQLConf.STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT.key -> "3"
-    ) {
-      withTempDir { checkpointDir =>
-        val inputData = MemoryStream[String]
-        val result = inputData.toDS().dropDuplicates()
-
-        testStream(result, outputMode = OutputMode.Update)(
-          StartStream(checkpointLocation = checkpointDir.getCanonicalPath),
-          AddData(inputData, "a"),
-          ProcessAllAvailable(),
-          AddData(inputData, "b"),
-          ProcessAllAvailable(),
-          CheckNewAnswer("a", "b"),
-          Execute { q =>
-            // Make sure only smallest K active metrics are published
-            eventually(timeout(10.seconds)) {
-              val instanceMetrics = q.lastProgress
-                .stateOperators(0)
-                .customMetrics
-                .asScala
-                .view
-                .filterKeys(_.startsWith(SNAPSHOT_LAG_METRIC_PREFIX))
-              // Determined by STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT
-              assert(
-                instanceMetrics.size == q.sparkSession.conf
-                  .get(SQLConf.STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT)
-              )
-              assert(instanceMetrics.forall(_._2 == 1))
-            }
-          },
-          StopStream
-        )
-      }
-    }
-  }
-
-  testWithChangelogCheckpointingEnabled(
-    "SPARK-51097: Verify snapshot lag metrics are updated correctly with " +
-    "SkipMaintenanceOnCertainPartitionsProvider"
-  ) {
-    withSQLConf(
-      SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
-        classOf[SkipMaintenanceOnCertainPartitionsProvider].getName,
-      SQLConf.STREAMING_MAINTENANCE_INTERVAL.key -> "100",
-      SQLConf.STREAMING_NO_DATA_PROGRESS_EVENT_INTERVAL.key -> "10",
-      SQLConf.STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT.key -> "1",
-      SQLConf.STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT.key -> "3"
-    ) {
-      withTempDir { checkpointDir =>
-        val inputData = MemoryStream[String]
-        val result = inputData.toDS().dropDuplicates()
-
-        testStream(result, outputMode = OutputMode.Update)(
-          StartStream(checkpointLocation = checkpointDir.getCanonicalPath),
-          AddData(inputData, "a"),
-          ProcessAllAvailable(),
-          AddData(inputData, "b"),
-          ProcessAllAvailable(),
-          CheckNewAnswer("a", "b"),
-          Execute { q =>
-            // Partitions getting skipped (id 0 and 1) do not have an uploaded version, leaving
-            // those instance metrics as -1.
-            eventually(timeout(10.seconds)) {
-              assert(
-                q.lastProgress
-                  .stateOperators(0)
-                  .customMetrics
-                  .get(snapshotLagMetricName(0)) === -1
-              )
-              assert(
-                q.lastProgress
-                  .stateOperators(0)
-                  .customMetrics
-                  .get(snapshotLagMetricName(1)) === -1
-              )
-              // Make sure only smallest K active metrics are published
-              val instanceMetrics = q.lastProgress
-                .stateOperators(0)
-                .customMetrics
-                .asScala
-                .view
-                .filterKeys(_.startsWith(SNAPSHOT_LAG_METRIC_PREFIX))
-              // Determined by STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT
-              assert(
-                instanceMetrics.size == q.sparkSession.conf
-                  .get(SQLConf.STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT)
-              )
-              // Two metrics published are -1, the remainder should all be 1 as they
-              // uploaded properly.
-              assert(
-                instanceMetrics.count(_._2 == 1) == q.sparkSession.conf
-                  .get(SQLConf.STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT) - 2
-              )
-            }
-          },
-          StopStream
-        )
-      }
-    }
-  }
-
-  testWithChangelogCheckpointingEnabled(
-    "SPARK-51097: Verify snapshot lag metrics are updated correctly for join queries with " +
-    "RocksDBStateStoreProvider"
-  ) {
-    withSQLConf(
-      SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
-        classOf[RocksDBStateStoreProvider].getName,
-      SQLConf.STREAMING_MAINTENANCE_INTERVAL.key -> "100",
-      SQLConf.STREAMING_NO_DATA_PROGRESS_EVENT_INTERVAL.key -> "10",
-      SQLConf.STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT.key -> "1",
-      SQLConf.STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT.key -> "10"
-    ) {
-      withTempDir { checkpointDir =>
-        val input1 = MemoryStream[Int]
-        val input2 = MemoryStream[Int]
-
-        val df1 = input1.toDF().select($"value" as "leftKey", ($"value" * 2) as "leftValue")
-        val df2 = input2
-          .toDF()
-          .select($"value" as "rightKey", ($"value" * 3) as "rightValue")
-        val joined = df1.join(df2, expr("leftKey = rightKey"))
-
-        testStream(joined)(
-          StartStream(checkpointLocation = checkpointDir.getCanonicalPath),
-          AddData(input1, 1, 5),
-          ProcessAllAvailable(),
-          AddData(input2, 1, 5, 10),
-          ProcessAllAvailable(),
-          CheckNewAnswer((1, 2, 1, 3), (5, 10, 5, 15)),
-          Execute { q =>
-            eventually(timeout(10.seconds)) {
-              // Make sure only smallest K active metrics are published.
-              // There are 5 * 4 = 20 metrics in total because of join, but only 10 are published.
-              val instanceMetrics = q.lastProgress
-                .stateOperators(0)
-                .customMetrics
-                .asScala
-                .view
-                .filterKeys(_.startsWith(SNAPSHOT_LAG_METRIC_PREFIX))
-              // Determined by STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT
-              assert(
-                instanceMetrics.size == q.sparkSession.conf
-                  .get(SQLConf.STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT)
-              )
-              // All state store instances should have uploaded a version
-              assert(instanceMetrics.forall(_._2 == 1))
-            }
-          },
-          StopStream
-        )
-      }
-    }
-  }
-
-  testWithChangelogCheckpointingEnabled(
-    "SPARK-51097: Verify snapshot lag metrics are updated correctly for join queries with " +
-    "SkipMaintenanceOnCertainPartitionsProvider"
-  ) {
-    withSQLConf(
-      SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
-        classOf[SkipMaintenanceOnCertainPartitionsProvider].getName,
-      SQLConf.STREAMING_MAINTENANCE_INTERVAL.key -> "100",
-      SQLConf.STREAMING_NO_DATA_PROGRESS_EVENT_INTERVAL.key -> "10",
-      SQLConf.STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT.key -> "1",
-      SQLConf.STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT.key -> "10"
-    ) {
-      withTempDir { checkpointDir =>
-        val input1 = MemoryStream[Int]
-        val input2 = MemoryStream[Int]
-
-        val df1 = input1.toDF().select($"value" as "leftKey", ($"value" * 2) as "leftValue")
-        val df2 = input2
-          .toDF()
-          .select($"value" as "rightKey", ($"value" * 3) as "rightValue")
-        val joined = df1.join(df2, expr("leftKey = rightKey"))
-
-        testStream(joined)(
-          StartStream(checkpointLocation = checkpointDir.getCanonicalPath),
-          AddData(input1, 1, 5),
-          ProcessAllAvailable(),
-          AddData(input2, 1, 5, 10),
-          ProcessAllAvailable(),
-          CheckNewAnswer((1, 2, 1, 3), (5, 10, 5, 15)),
-          Execute { q =>
-            eventually(timeout(10.seconds)) {
-              // Make sure only smallest K active metrics are published.
-              // There are 5 * 4 = 20 metrics in total because of join, but only 10 are published.
-              val allInstanceMetrics = q.lastProgress
-                .stateOperators(0)
-                .customMetrics
-                .asScala
-                .view
-                .filterKeys(_.startsWith(SNAPSHOT_LAG_METRIC_PREFIX))
-              val badInstanceMetrics = allInstanceMetrics.filterKeys(
-                k =>
-                  k.startsWith(snapshotLagMetricName(0, "")) ||
-                  k.startsWith(snapshotLagMetricName(1, ""))
-              )
-              // Determined by STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT
-              assert(
-                allInstanceMetrics.size == q.sparkSession.conf
-                  .get(SQLConf.STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT)
-              )
-              // Two ids are blocked, each with four state stores
-              assert(badInstanceMetrics.count(_._2 == -1) == 2 * 4)
-              // The rest should have uploaded a version
-              assert(
-                allInstanceMetrics.count(_._2 == 1) == q.sparkSession.conf
-                  .get(SQLConf.STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT) - 2 * 4
-              )
-            }
-          },
-          StopStream
-        )
-      }
-    }
-  }
-
-  testWithChangelogCheckpointingEnabled(
-    "SPARK-51097: Verify RocksDB instance metrics are not collected in execution plan"
-  ) {
-    withSQLConf(
-      SQLConf.STATE_STORE_PROVIDER_CLASS.key -> classOf[RocksDBStateStoreProvider].getName
-    ) {
-      withTempDir { checkpointDir =>
-        val input1 = MemoryStream[Int]
-        val input2 = MemoryStream[Int]
-
-        val df1 = input1.toDF().select($"value" as "leftKey", ($"value" * 2) as "leftValue")
-        val df2 = input2
-          .toDF()
-          .select($"value" as "rightKey", ($"value" * 3) as "rightValue")
-        val joined = df1.join(df2, expr("leftKey = rightKey"))
-
-        testStream(joined)(
-          StartStream(checkpointLocation = checkpointDir.getCanonicalPath),
-          AddData(input1, 1, 5),
-          ProcessAllAvailable(),
-          AddData(input2, 1, 5, 10),
-          ProcessAllAvailable(),
-          CheckNewAnswer((1, 2, 1, 3), (5, 10, 5, 15)),
-          AssertOnQuery { q =>
-            // Go through all elements in the execution plan and verify none of the metrics
-            // are generated from RocksDB's snapshot lag instance metrics.
-            q.lastExecution.executedPlan
-              .collect {
-                case node => node.metrics
-              }
-              .forall { nodeMetrics =>
-                nodeMetrics.forall(metric => !metric._1.startsWith(SNAPSHOT_LAG_METRIC_PREFIX))
-              }
-          },
-          StopStream
-        )
-      }
-    }
-  }
 }