[SPARK-41054][UI][CORE] Support RocksDB as KVStore in live UI

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

Support using RocksDB as the KVStore in live UI. The location of RocksDB can be set via configuration `spark.ui.store.path`. The configuration is optional. The default KV store will still be in memory.

Note: let's make it ROCKSDB only for now instead of supporting both LevelDB and RocksDB. RocksDB is built based on LevelDB with improvements on writes and reads. Furthermore, we can reuse the RocksDBFileManager in streaming for replicating the local RocksDB file to DFS. The replication in DFS can be used for the Spark history server.

### Why are the changes needed?

The current architecture of Spark live UI and Spark history server(SHS) is too simple to serve large clusters and heavy workloads:

- Spark stores all the live UI date in memory. The size can be a few GBs and affects the driver's stability (OOM).
- There is a limitation of storing 1000 queries only. Note that we can’t simply increase the limitation under the current Architecture. I did a memory profiling. Storing one query execution detail can take 800KB while storing one task requires 0.3KB. So for 1000 SQL queries with 1000* 2000 tasks, the memory usage for query execution and task data will be 1.4GB. Spark UI stores UI data for jobs/stages/executors as well.  So to store 10k queries, it may take more than 14GB.
- SHS has to parse JSON format event log for the initial start.  The uncompressed event logs can be as big as a few GBs, and the parse can be quite slow. Some users reported they had to wait for more than half an hour.

With RocksDB as KVStore, we can improve the stability of Spark driver and fasten the startup of SHS.

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

Yes, supporting RocksDB as the KVStore in live UI. The location of RocksDB can be set via configuration `spark.ui.store.path`. The configuration is optional. The default KV store will still be in memory.

### How was this patch tested?

New UT
Preview of the doc change:
<img width="895" alt="image" src="https://user-images.githubusercontent.com/1097932/203184691-b6815990-b7b0-422b-aded-8e1771c0c167.png">

Closes #38567 from gengliangwang/liveUIKVStore.

Authored-by: Gengliang Wang <[email protected]>
Signed-off-by: Gengliang Wang <[email protected]>

Author: gengliangwang

core/src/main/scala/org/apache/spark/deploy/history/FsHistoryProvider.scala

@@ -19,7 +19,6 @@ package org.apache.spark.deploy.history
 
 import java.io.{File, FileNotFoundException, IOException}
 import java.lang.{Long => JLong}
-import java.nio.file.Files
 import java.util.{Date, NoSuchElementException, ServiceLoader}
 import java.util.concurrent.{ConcurrentHashMap, ExecutorService, TimeUnit}
 import java.util.zip.ZipOutputStream
@@ -36,15 +35,12 @@ import org.apache.hadoop.fs.{FileStatus, FileSystem, Path}
 import org.apache.hadoop.hdfs.DistributedFileSystem
 import org.apache.hadoop.hdfs.protocol.HdfsConstants
 import org.apache.hadoop.security.AccessControlException
-import org.fusesource.leveldbjni.internal.NativeDB
-import org.rocksdb.RocksDBException
 
 import org.apache.spark.{SecurityManager, SparkConf, SparkException}
 import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.internal.Logging
 import org.apache.spark.internal.config._
 import org.apache.spark.internal.config.History._
-import org.apache.spark.internal.config.History.HybridStoreDiskBackend._
 import org.apache.spark.internal.config.Status._
 import org.apache.spark.internal.config.Tests.IS_TESTING
 import org.apache.spark.internal.config.UI._
@@ -136,34 +132,9 @@ private[history] class FsHistoryProvider(conf: SparkConf, clock: Clock)
     HybridStoreDiskBackend.withName(conf.get(History.HYBRID_STORE_DISK_BACKEND))
 
   // Visible for testing.
-  private[history] val listing: KVStore = storePath.map { path =>
-    val dir = hybridStoreDiskBackend match {
-      case LEVELDB => "listing.ldb"
-      case ROCKSDB => "listing.rdb"
-    }
-    val dbPath = Files.createDirectories(new File(path, dir).toPath()).toFile()
-    Utils.chmod700(dbPath)
-
-    val metadata = FsHistoryProviderMetadata(CURRENT_LISTING_VERSION,
-      AppStatusStore.CURRENT_VERSION, logDir)
-
-    try {
-      open(dbPath, metadata, conf)
-    } catch {
-      // If there's an error, remove the listing database and any existing UI database
-      // from the store directory, since it's extremely likely that they'll all contain
-      // incompatible information.
-      case _: UnsupportedStoreVersionException | _: MetadataMismatchException =>
-        logInfo("Detected incompatible DB versions, deleting...")
-        path.listFiles().foreach(Utils.deleteRecursively)
-        open(dbPath, metadata, conf)
-      case dbExc @ (_: NativeDB.DBException | _: RocksDBException) =>
-        // Get rid of the corrupted data and re-create it.
-        logWarning(s"Failed to load disk store $dbPath :", dbExc)
-        Utils.deleteRecursively(dbPath)
-        open(dbPath, metadata, conf)
-    }
-  }.getOrElse(new InMemoryStore())
+  private[history] val listing: KVStore = {
+    KVUtils.createKVStore(storePath, hybridStoreDiskBackend, conf)
+  }
 
   private val diskManager = storePath.map { path =>
     new HistoryServerDiskManager(conf, path, listing, clock)
@@ -1444,7 +1415,7 @@ private[history] class FsHistoryProvider(conf: SparkConf, clock: Clock)
   }
 }
 
-private[history] object FsHistoryProvider {
+private[spark] object FsHistoryProvider {
 
   private val APPL_START_EVENT_PREFIX = "{\"Event\":\"SparkListenerApplicationStart\""
 
@@ -1459,10 +1430,10 @@ private[history] object FsHistoryProvider {
    * db, if the version does not match this value, the FsHistoryProvider will throw away
    * all data and re-generate the listing data from the event logs.
    */
-  private[history] val CURRENT_LISTING_VERSION = 1L
+  val CURRENT_LISTING_VERSION = 1L
 }
 
-private[history] case class FsHistoryProviderMetadata(
+private[spark] case class FsHistoryProviderMetadata(
     version: Long,
     uiVersion: Long,
     logDir: String)

core/src/main/scala/org/apache/spark/internal/config/Status.scala

@@ -70,4 +70,11 @@ private[spark] object Status {
       .version("3.0.0")
       .booleanConf
       .createWithDefault(false)
+
+  val LIVE_UI_LOCAL_STORE_DIR = ConfigBuilder("spark.ui.store.path")
+    .doc("Local directory where to cache application information for live UI. By default this is " +
+      "not set, meaning all application information will be kept in memory.")
+    .version("3.4.0")
+    .stringConf
+    .createOptional
 }

core/src/main/scala/org/apache/spark/status/AppStatusStore.scala

@@ -17,17 +17,20 @@
 
 package org.apache.spark.status
 
+import java.io.File
 import java.util.{List => JList}
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.HashMap
 
 import org.apache.spark.{JobExecutionStatus, SparkConf, SparkContext}
+import org.apache.spark.internal.config.History.HybridStoreDiskBackend
+import org.apache.spark.internal.config.Status.LIVE_UI_LOCAL_STORE_DIR
 import org.apache.spark.status.api.v1
 import org.apache.spark.storage.FallbackStorage.FALLBACK_BLOCK_MANAGER_ID
 import org.apache.spark.ui.scope._
 import org.apache.spark.util.Utils
-import org.apache.spark.util.kvstore.{InMemoryStore, KVStore}
+import org.apache.spark.util.kvstore.KVStore
 
 /**
  * A wrapper around a KVStore that provides methods for accessing the API data stored within.
@@ -769,7 +772,12 @@ private[spark] object AppStatusStore {
   def createLiveStore(
       conf: SparkConf,
       appStatusSource: Option[AppStatusSource] = None): AppStatusStore = {
-    val store = new ElementTrackingStore(new InMemoryStore(), conf)
+    val storePath = conf.get(LIVE_UI_LOCAL_STORE_DIR).map(new File(_))
+    // For the disk-based KV store of live UI, let's simply make it ROCKSDB only for now,
+    // instead of supporting both LevelDB and RocksDB. RocksDB is built based on LevelDB with
+    // improvements on writes and reads.
+    val kvStore = KVUtils.createKVStore(storePath, HybridStoreDiskBackend.ROCKSDB, conf)
+    val store = new ElementTrackingStore(kvStore, conf)
     val listener = new AppStatusListener(store, conf, true, appStatusSource)
     new AppStatusStore(store, listener = Some(listener))
   }

core/src/main/scala/org/apache/spark/status/KVUtils.scala

@@ -18,16 +18,21 @@
 package org.apache.spark.status
 
 import java.io.File
+import java.nio.file.Files
 
 import scala.annotation.meta.getter
 import scala.collection.JavaConverters._
 import scala.reflect.{classTag, ClassTag}
 
 import com.fasterxml.jackson.annotation.JsonInclude
 import com.fasterxml.jackson.module.scala.DefaultScalaModule
+import org.fusesource.leveldbjni.internal.NativeDB
+import org.rocksdb.RocksDBException
 
 import org.apache.spark.SparkConf
+import org.apache.spark.deploy.history.{FsHistoryProvider, FsHistoryProviderMetadata}
 import org.apache.spark.internal.Logging
+import org.apache.spark.internal.config.History
 import org.apache.spark.internal.config.History.HYBRID_STORE_DISK_BACKEND
 import org.apache.spark.internal.config.History.HybridStoreDiskBackend
 import org.apache.spark.internal.config.History.HybridStoreDiskBackend._
@@ -62,10 +67,14 @@ private[spark] object KVUtils extends Logging {
    *                 the store's metadata.
    * @param conf SparkConf use to get `HYBRID_STORE_DISK_BACKEND`
    */
-  def open[M: ClassTag](path: File, metadata: M, conf: SparkConf): KVStore = {
+  def open[M: ClassTag](
+      path: File,
+      metadata: M,
+      conf: SparkConf,
+      diskBackend: Option[HybridStoreDiskBackend.Value] = None): KVStore = {
     require(metadata != null, "Metadata is required.")
 
-    val db = backend(conf) match {
+    val db = diskBackend.getOrElse(backend(conf)) match {
       case LEVELDB => new LevelDB(path, new KVStoreScalaSerializer())
       case ROCKSDB => new RocksDB(path, new KVStoreScalaSerializer())
     }
@@ -80,6 +89,43 @@ private[spark] object KVUtils extends Logging {
     db
   }
 
+  def createKVStore(
+      storePath: Option[File],
+      diskBackend: HybridStoreDiskBackend.Value,
+      conf: SparkConf): KVStore = {
+    storePath.map { path =>
+      val dir = diskBackend match {
+        case LEVELDB => "listing.ldb"
+        case ROCKSDB => "listing.rdb"
+      }
+
+      val dbPath = Files.createDirectories(new File(path, dir).toPath()).toFile()
+      Utils.chmod700(dbPath)
+
+      val metadata = FsHistoryProviderMetadata(
+        FsHistoryProvider.CURRENT_LISTING_VERSION,
+        AppStatusStore.CURRENT_VERSION,
+        conf.get(History.HISTORY_LOG_DIR))
+
+      try {
+        open(dbPath, metadata, conf, Some(diskBackend))
+      } catch {
+        // If there's an error, remove the listing database and any existing UI database
+        // from the store directory, since it's extremely likely that they'll all contain
+        // incompatible information.
+        case _: UnsupportedStoreVersionException | _: MetadataMismatchException =>
+          logInfo("Detected incompatible DB versions, deleting...")
+          path.listFiles().foreach(Utils.deleteRecursively)
+          open(dbPath, metadata, conf, Some(diskBackend))
+        case dbExc @ (_: NativeDB.DBException | _: RocksDBException) =>
+          // Get rid of the corrupted data and re-create it.
+          logWarning(s"Failed to load disk store $dbPath :", dbExc)
+          Utils.deleteRecursively(dbPath)
+          open(dbPath, metadata, conf, Some(diskBackend))
+      }
+    }.getOrElse(new InMemoryStore())
+  }
+
   /** Turns a KVStoreView into a Scala sequence, applying a filter. */
   def viewToSeq[T](
       view: KVStoreView[T],

core/src/test/scala/org/apache/spark/status/AppStatusStoreSuite.scala

@@ -22,7 +22,7 @@ import scala.util.Random
 import org.apache.spark.{SparkConf, SparkFunSuite}
 import org.apache.spark.executor.TaskMetrics
 import org.apache.spark.internal.config.History.{HYBRID_STORE_DISK_BACKEND, HybridStoreDiskBackend}
-import org.apache.spark.internal.config.Status.LIVE_ENTITY_UPDATE_PERIOD
+import org.apache.spark.internal.config.Status.{LIVE_ENTITY_UPDATE_PERIOD, LIVE_UI_LOCAL_STORE_DIR}
 import org.apache.spark.resource.ResourceProfile
 import org.apache.spark.scheduler.{SparkListenerStageSubmitted, SparkListenerTaskStart, StageInfo, TaskInfo, TaskLocality}
 import org.apache.spark.status.api.v1.SpeculationStageSummary
@@ -88,7 +88,19 @@ class AppStatusStoreSuite extends SparkFunSuite {
       live: Boolean): AppStatusStore = {
     val conf = new SparkConf()
     if (live) {
-      return AppStatusStore.createLiveStore(conf)
+      if (disk) {
+        val testDir = Utils.createTempDir()
+        conf.set(LIVE_UI_LOCAL_STORE_DIR, testDir.getCanonicalPath)
+      }
+      val liveStore = AppStatusStore.createLiveStore(conf)
+      if (disk) {
+        val rocksDBCreated = liveStore.store match {
+          case e: ElementTrackingStore => !e.usingInMemoryStore
+          case _ => false
+        }
+        assert(rocksDBCreated)
+      }
+      return liveStore
     }
 
     val store: KVStore = if (disk) {
@@ -106,7 +118,8 @@ class AppStatusStoreSuite extends SparkFunSuite {
     val baseCases = Seq(
       "disk rocksdb" -> createAppStore(disk = true, HybridStoreDiskBackend.ROCKSDB, live = false),
       "in memory" -> createAppStore(disk = false, live = false),
-      "in memory live" -> createAppStore(disk = false, live = true)
+      "in memory live" -> createAppStore(disk = false, live = true),
+      "rocksdb live" -> createAppStore(disk = true, HybridStoreDiskBackend.ROCKSDB, live = true)
     )
     if (Utils.isMacOnAppleSilicon) {
       baseCases

docs/configuration.md

@@ -1379,6 +1379,15 @@ Apart from these, the following properties are also available, and may be useful
   </td>
   <td>1.1.1</td>
 </tr>
+<tr>
+  <td><code>spark.ui.store.path</code></td>
+  <td>None</td>
+  <td>
+    Local directory where to cache application information for live UI.
+    By default this is not set, meaning all application information will be kept in memory.
+  </td>
+  <td>3.4.0</td>
+</tr>
 <tr>
   <td><code>spark.ui.killEnabled</code></td>
   <td>true</td>