SPARK-51252 Switch to atomic and fix spark test prefix

zecookiez · Feb 24, 2025 · 10cbd2e · 10cbd2e
1 parent e219d2b
commit 10cbd2e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 11 deletions.
diff --git a/...n/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala b/...n/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._
@@ -221,8 +221,7 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
 
       val instanceMetrics = Map(
         instanceMetricSnapshotLastUpload.withNewId(
-          stateStoreId.partitionId, stateStoreId.storeName
-        ) -> lastSnapshotUploadedVersion
+          stateStoreId.partitionId, stateStoreId.storeName) -> lastSnapshotUploadedVersion.get()
       )
 
       StateStoreMetrics(
@@ -433,9 +432,9 @@ 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, -1 represents no version
-  // has ever been uploaded.
-  private var lastSnapshotUploadedVersion: Long = -1L
+  // 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 lastSnapshotUploadedVersion: AtomicLong = new AtomicLong(-1L)
 
   private lazy val metricStateOnCurrentVersionSizeBytes: StateStoreCustomSizeMetric =
     StateStoreCustomSizeMetric("stateOnCurrentVersionSizeBytes",
@@ -698,7 +697,7 @@ 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)}")
-    lastSnapshotUploadedVersion = version
+    lastSnapshotUploadedVersion.set(version)
   }
 
   /**

diff --git a/.../scala/org/apache/spark/sql/execution/streaming/state/StateStoreInstanceMetricSuite.scala b/.../scala/org/apache/spark/sql/execution/streaming/state/StateStoreInstanceMetricSuite.scala
@@ -62,7 +62,7 @@ class StateStoreInstanceMetricSuite extends StreamTest with AlsoTestWithRocksDBF
 
   Seq(
     ("SPARK-51097", "RocksDBStateStoreProvider", classOf[RocksDBStateStoreProvider].getName),
-    ("SPARK-51251", "HDFSBackedStateStoreProvider", classOf[HDFSBackedStateStoreProvider].getName)
+    ("SPARK-51252", "HDFSBackedStateStoreProvider", classOf[HDFSBackedStateStoreProvider].getName)
   ).foreach {
     case (ticketPrefix, providerName, providerClassName) =>
       testWithChangelogCheckpointingEnabled(
@@ -120,7 +120,7 @@ class StateStoreInstanceMetricSuite extends StreamTest with AlsoTestWithRocksDBF
       classOf[RocksDBSkipMaintenanceOnCertainPartitionsProvider].getName
     ),
     (
-      "SPARK-51251",
+      "SPARK-51252",
       "HDFSBackedSkipMaintenanceOnCertainPartitionsProvider",
       classOf[HDFSBackedSkipMaintenanceOnCertainPartitionsProvider].getName
     )
@@ -194,7 +194,7 @@ class StateStoreInstanceMetricSuite extends StreamTest with AlsoTestWithRocksDBF
 
   Seq(
     ("SPARK-51097", "RocksDBStateStoreProvider", classOf[RocksDBStateStoreProvider].getName),
-    ("SPARK-51251", "HDFSBackedStateStoreProvider", classOf[HDFSBackedStateStoreProvider].getName)
+    ("SPARK-51252", "HDFSBackedStateStoreProvider", classOf[HDFSBackedStateStoreProvider].getName)
   ).foreach {
     case (ticketPrefix, providerName, providerClassName) =>
       testWithChangelogCheckpointingEnabled(
@@ -261,7 +261,7 @@ class StateStoreInstanceMetricSuite extends StreamTest with AlsoTestWithRocksDBF
       classOf[RocksDBSkipMaintenanceOnCertainPartitionsProvider].getName
     ),
     (
-      "SPARK-51251",
+      "SPARK-51252",
       "HDFSBackedSkipMaintenanceOnCertainPartitionsProvider",
       classOf[HDFSBackedSkipMaintenanceOnCertainPartitionsProvider].getName
     )