[SPARK-43421][SS] Implement Changelog based Checkpointing for RocksDB State Store Provider

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

@@ -1940,6 +1940,15 @@ object SQLConf {
       // 5 is the default table format version for RocksDB 6.20.3.
       .createWithDefault(5)
 
+  val STATE_STORE_ROCKSDB_CHANGE_CHECKPOINTING_ENABLED =
+    buildConf("spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled")
+      .internal()
+      .doc("Enable RocksDB state store to checkpoint a version of the state" +
+        " by uploading the changelog.")
+      .version("3.4.1")
+      .booleanConf
+      .createWithDefault(false)
+
   val STREAMING_AGGREGATION_STATE_FORMAT_VERSION =
     buildConf("spark.sql.streaming.aggregation.stateFormatVersion")
       .internal()

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

@@ -56,6 +56,15 @@ class RocksDB(
     hadoopConf: Configuration = new Configuration,
     loggingId: String = "") extends Logging {
 
+  case class RocksDBCheckpoint(checkpointDir: File, version: Long, numKeys: Long) {
+    def close(): Unit = {
+      silentDeleteRecursively(checkpointDir, s"Free up local checkpoint of snapshot $version")
+    }
+  }
+
+  @volatile private var latestCheckpoint: Option[RocksDBCheckpoint] = None
+  @volatile private var lastCheckpointVersion = 0L
+
   RocksDBLoader.loadLibrary()
 
   // Java wrapper objects linking to native RocksDB objects
@@ -109,13 +118,15 @@ class RocksDB(
   private val nativeStats = dbOptions.statistics()
 
   private val workingDir = createTempDir("workingDir")
-  private val fileManager = new RocksDBFileManager(
-    dfsRootDir, createTempDir("fileManager"), hadoopConf, loggingId = loggingId)
+  private val fileManager = new RocksDBFileManager(dfsRootDir, createTempDir("fileManager"),
+    hadoopConf, conf.compressionCodec, loggingId = loggingId)
   private val byteArrayPair = new ByteArrayPair()
   private val commitLatencyMs = new mutable.HashMap[String, Long]()
   private val acquireLock = new Object
 
   @volatile private var db: NativeRocksDB = _
+  @volatile private var changelogWriter: Option[StateStoreChangelogWriter] = None
+  private val enableChangelogCheckpointing: Boolean = conf.enableChangelogCheckpointing
   @volatile private var loadedVersion = -1L   // -1 = nothing valid is loaded
   @volatile private var numKeysOnLoadedVersion = 0L
   @volatile private var numKeysOnWritingVersion = 0L
@@ -129,14 +140,18 @@ class RocksDB(
    * Note that this will copy all the necessary file from DFS to local disk as needed,
    * and possibly restart the native RocksDB instance.
    */
-  def load(version: Long): RocksDB = {
+  def load(version: Long, readOnly: Boolean = false): RocksDB = {
     assert(version >= 0)
     acquire()
     logInfo(s"Loading $version")
     try {
       if (loadedVersion != version) {
         closeDB()
-        val metadata = fileManager.loadCheckpointFromDfs(version, workingDir)
+        val loadSnapshotVersion = if (enableChangelogCheckpointing) {
+          fileManager.getLatestSnapshotVersion(version)
+        } else version
+        val metadata = fileManager.loadCheckpointFromDfs(loadSnapshotVersion, workingDir)
+        loadedVersion = loadSnapshotVersion
         openDB()
 
         val numKeys = if (!conf.trackTotalNumberOfRows) {
@@ -150,9 +165,10 @@ class RocksDB(
           metadata.numKeys
         }
         numKeysOnWritingVersion = numKeys
-        numKeysOnLoadedVersion = numKeys
-
-        loadedVersion = version
+        if (enableChangelogCheckpointing) replayChangelog(version)
+        // After changelog replay the numKeysOnWritingVersion will be updated to
+        // the correct number of keys in the loaded version.
+        numKeysOnLoadedVersion = numKeysOnWritingVersion
         fileManagerMetrics = fileManager.latestLoadCheckpointMetrics
       }
       if (conf.resetStatsOnLoad) {
@@ -164,9 +180,34 @@ class RocksDB(
         loadedVersion = -1  // invalidate loaded data
         throw t
     }
+    if (enableChangelogCheckpointing && !readOnly) {
+      changelogWriter = Some(fileManager.getChangeLogWriter(version + 1))
+    }
     this
   }
 
+  private def replayChangelog(endVersion: Long): Unit = {
+    // Replay change log from the loaded version to the target version.
+    // This will be noop if changelog checkpointing is disabled.
+    for (v <- loadedVersion + 1 to endVersion) {
+      var changelogReader: StateStoreChangelogReader = null
+      try {
+        changelogReader = fileManager.getChangelogReader(v)
+        while (changelogReader.hasNext) {
+          val byteArrayPair = changelogReader.next()
+          if (byteArrayPair.value != null) {
+            put(byteArrayPair.key, byteArrayPair.value)
+          } else {
+            remove(byteArrayPair.key)
+          }
+        }
+      } finally {
+        if (changelogReader != null) changelogReader.close()
+      }
+    }
+    loadedVersion = endVersion
+  }
+
   /**
    * Get the value for the given key if present, or null.
    * @note This will return the last written value even if it was uncommitted.
@@ -187,6 +228,7 @@ class RocksDB(
       }
     }
     db.put(writeOptions, key, value)
+    changelogWriter.foreach(_.put(key, value))
   }
 
   /**
@@ -201,6 +243,7 @@ class RocksDB(
       }
     }
     db.delete(writeOptions, key)
+    changelogWriter.foreach(_.delete(key))
   }
 
   /**
@@ -286,44 +329,51 @@ class RocksDB(
    */
   def commit(): Long = {
     val newVersion = loadedVersion + 1
-    val checkpointDir = createTempDir("checkpoint")
-    var rocksDBBackgroundThreadPaused = false
     try {
-      // Make sure the directory does not exist. Native RocksDB fails if the directory to
-      // checkpoint exists.
-      Utils.deleteRecursively(checkpointDir)
 
       logInfo(s"Flushing updates for $newVersion")
-      val flushTimeMs = timeTakenMs { db.flush(flushOptions) }
-
-      val compactTimeMs = if (conf.compactOnCommit) {
-        logInfo("Compacting")
-        timeTakenMs { db.compactRange() }
-      } else 0
-
-      logInfo("Pausing background work")
-      val pauseTimeMs = timeTakenMs {
-        db.pauseBackgroundWork() // To avoid files being changed while committing
-        rocksDBBackgroundThreadPaused = true
-      }
 
-      logInfo(s"Creating checkpoint for $newVersion in $checkpointDir")
-      val checkpointTimeMs = timeTakenMs {
-        val cp = Checkpoint.create(db)
-        cp.createCheckpoint(checkpointDir.toString)
+      var flushTimeMs = 0L
+      var checkpointTimeMs = 0L
+      if (shouldCreateSnapshot()) {
+        // Need to flush the change to disk before creating a checkpoint
+        // because rocksdb wal is disabled.
+        flushTimeMs = timeTakenMs { db.flush(flushOptions) }
+        checkpointTimeMs = timeTakenMs {
+          val checkpointDir = createTempDir("checkpoint")
+          // Make sure the directory does not exist. Native RocksDB fails if the directory to
+          // checkpoint exists.
+          Utils.deleteRecursively(checkpointDir)
+          val cp = Checkpoint.create(db)
+          cp.createCheckpoint(checkpointDir.toString)
+          synchronized {
+            latestCheckpoint.foreach(_.close())
+            latestCheckpoint = Some(
+              RocksDBCheckpoint(checkpointDir, newVersion, numKeysOnWritingVersion))
+            lastCheckpointVersion = newVersion
+          }
+        }
       }
 
       logInfo(s"Syncing checkpoint for $newVersion to DFS")
       val fileSyncTimeMs = timeTakenMs {
-        fileManager.saveCheckpointToDfs(checkpointDir, newVersion, numKeysOnWritingVersion)
+        if (enableChangelogCheckpointing) {
+          try {
+            assert(changelogWriter.isDefined)
+            changelogWriter.foreach(_.commit())
+          } finally {
+            changelogWriter = None
+          }
+        } else {
+          assert(changelogWriter.isEmpty)
+          uploadSnapshot()
+        }
       }
+
       numKeysOnLoadedVersion = numKeysOnWritingVersion
       loadedVersion = newVersion
-      fileManagerMetrics = fileManager.latestSaveCheckpointMetrics
       commitLatencyMs ++= Map(
         "flush" -> flushTimeMs,
-        "compact" -> compactTimeMs,
-        "pause" -> pauseTimeMs,
         "checkpoint" -> checkpointTimeMs,
         "fileSync" -> fileSyncTimeMs
       )
@@ -334,25 +384,59 @@ class RocksDB(
         loadedVersion = -1  // invalidate loaded version
         throw t
     } finally {
-      if (rocksDBBackgroundThreadPaused) db.continueBackgroundWork()
-      silentDeleteRecursively(checkpointDir, s"committing $newVersion")
       // reset resources as either 1) we already pushed the changes and it has been committed or
       // 2) commit has failed and the current version is "invalidated".
       release()
     }
   }
 
+  private def shouldCreateSnapshot(): Boolean = {
+    if (enableChangelogCheckpointing) {
+      assert(changelogWriter.isDefined)
+      val newVersion = loadedVersion + 1
+      newVersion - lastCheckpointVersion >= conf.minDeltasForSnapshot ||
+        changelogWriter.get.size > 10000
+    } else true
+  }
+
+  private def uploadSnapshot(): Unit = {
+    val localCheckpoint = synchronized {
+      val checkpoint = latestCheckpoint
+      latestCheckpoint = None
+      checkpoint
+    }
+    localCheckpoint match {
+      case Some(RocksDBCheckpoint(localDir, version, numKeys)) =>
+        try {
+          val uploadTime = timeTakenMs {
+            fileManager.saveCheckpointToDfs(localDir, version, numKeys)
+            fileManagerMetrics = fileManager.latestSaveCheckpointMetrics
+          }
+          logInfo(s"Upload snapshot of version $version, time taken: $uploadTime ms")
+        } finally {
+          localCheckpoint.foreach(_.close())
+        }
+      case _ =>
+    }
+  }
+
   /**
    * Drop uncommitted changes, and roll back to previous version.
    */
   def rollback(): Unit = {
     numKeysOnWritingVersion = numKeysOnLoadedVersion
     loadedVersion = -1L
+    changelogWriter.foreach(_.abort())
+    // Make sure changelogWriter gets recreated next time.
+    changelogWriter = None
     release()
     logInfo(s"Rolled back to $loadedVersion")
   }
 
-  def cleanup(): Unit = {
+  def doMaintenance(): Unit = {
+    if (enableChangelogCheckpointing) {
+      uploadSnapshot()
+    }
     val cleanupTime = timeTakenMs {
       fileManager.deleteOldVersions(conf.minVersionsToRetain)
     }
@@ -369,6 +453,7 @@ class RocksDB(
       flushOptions.close()
       dbOptions.close()
       dbLogger.close()
+      latestCheckpoint.foreach(_.close())
       silentDeleteRecursively(localRootDir, "closing RocksDB")
     } catch {
       case e: Exception =>
@@ -550,7 +635,9 @@ class ByteArrayPair(var key: Array[Byte] = null, var value: Array[Byte] = null)
  */
 case class RocksDBConf(
     minVersionsToRetain: Int,
+    minDeltasForSnapshot: Int,
     compactOnCommit: Boolean,
+    enableChangelogCheckpointing: Boolean,
     blockSizeKB: Long,
     blockCacheSizeMB: Long,
     lockAcquireTimeoutMs: Long,
@@ -563,7 +650,8 @@ case class RocksDBConf(
     boundedMemoryUsage: Boolean,
     totalMemoryUsageMB: Long,
     writeBufferCacheRatio: Double,
-    highPriorityPoolRatio: Double)
+    highPriorityPoolRatio: Double,
+    compressionCodec: String)
 
 object RocksDBConf {
   /** Common prefix of all confs in SQLConf that affects RocksDB */
@@ -585,6 +673,8 @@ object RocksDBConf {
 
   // Configuration that specifies whether to compact the RocksDB data every time data is committed
   private val COMPACT_ON_COMMIT_CONF = SQLConfEntry("compactOnCommit", "false")
+  private val ENABLE_CHANGELOG_CHECKPOINTING_CONF = SQLConfEntry(
+    "enableChangelogCheckpointing", "false")
   private val BLOCK_SIZE_KB_CONF = SQLConfEntry("blockSizeKB", "4")
   private val BLOCK_CACHE_SIZE_MB_CONF = SQLConfEntry("blockCacheSizeMB", "8")
   // See SPARK-42794 for details.
@@ -705,7 +795,9 @@ object RocksDBConf {
 
     RocksDBConf(
       storeConf.minVersionsToRetain,
+      storeConf.minDeltasForSnapshot,
       getBooleanConf(COMPACT_ON_COMMIT_CONF),
+      getBooleanConf(ENABLE_CHANGELOG_CHECKPOINTING_CONF),
       getPositiveLongConf(BLOCK_SIZE_KB_CONF),
       getPositiveLongConf(BLOCK_CACHE_SIZE_MB_CONF),
       getPositiveLongConf(LOCK_ACQUIRE_TIMEOUT_MS_CONF),
@@ -718,7 +810,8 @@ object RocksDBConf {
       getBooleanConf(BOUNDED_MEMORY_USAGE_CONF),
       getLongConf(MAX_MEMORY_USAGE_MB_CONF),
       getRatioConf(WRITE_BUFFER_CACHE_RATIO_CONF),
-      getRatioConf(HIGH_PRIORITY_POOL_RATIO_CONF))
+      getRatioConf(HIGH_PRIORITY_POOL_RATIO_CONF),
+      storeConf.compressionCodec)
   }
 
   def apply(): RocksDBConf = apply(new StateStoreConf())