[SPARK-53792] Fix rocksdbPinnedBlocksMemoryUsage when bounded memory usage is enabled

Author: Ubuntu

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

@@ -1461,7 +1461,6 @@ class RocksDB(
   private def metrics: RocksDBMetrics = {
     import HistogramType._
     val totalSSTFilesBytes = getDBProperty("rocksdb.total-sst-files-size")
-    val pinnedBlocksMemUsage = getDBProperty("rocksdb.block-cache-pinned-usage")
     val nativeOpsHistograms = Seq(
       "get" -> DB_GET,
       "put" -> DB_WRITE,
@@ -1497,6 +1496,11 @@ class RocksDB(
     // Use RocksDBMemoryManager to calculate the memory usage accounting
     val memoryUsage = RocksDBMemoryManager.getInstanceMemoryUsage(instanceUniqueId, getMemoryUsage)
 
+    val requestedPinnedBlocksMemUsage = getDBProperty("rocksdb.block-cache-pinned-usage")
+    val globalPinnedBlocksMemUsage = lruCache.getPinnedUsage()
+    val pinnedBlocksMemUsage = RocksDBMemoryManager.getInstancePinnedBlocksMemUsage(
+      instanceUniqueId, globalPinnedBlocksMemUsage, requestedPinnedBlocksMemUsage)
+
     RocksDBMetrics(
       numKeysOnLoadedVersion,
       numKeysOnWritingVersion,

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

@@ -117,10 +117,45 @@ object RocksDBMemoryManager extends Logging with UnmanagedMemoryConsumer {
       totalMemoryUsage / numBoundedInstances
     } else {
       // In unbounded memory mode, each instance has its own memory
+
+      // Question: Should we return 0L if the instance is not registed in instanceMemoryMap?
+      // but when it is not registed?
       totalMemoryUsage
     }
   }
 
+  /**
+   * Get the pinned blocks memory usage for a specific instance, accounting for bounded memory
+   * sharing.
+   * @param uniqueId The instance's unique identifier
+   * @param globalPinnedUsage The total pinned usage from the cache
+   * @return The adjusted pinned blocks memory usage accounting for sharing in bounded memory mode
+   */
+  def getInstancePinnedBlocksMemUsage(
+      uniqueId: String,
+      globalPinnedUsage: Long,
+      requestedPinnedUsage: Long): Long = {
+    if (!instanceMemoryMap.containsKey(uniqueId)) {
+      // Instance not registered, return 0
+      return 0L
+    }
+
+    val instanceInfo = instanceMemoryMap.get(uniqueId)
+    if (instanceInfo.isBoundedMemory) {
+      // In bounded memory mode, divide by the number of bounded instances
+      // since they share the same cache
+      val numBoundedInstances = getNumRocksDBInstances(true /* boundedMemory */)
+      if (numBoundedInstances > 0) {
+        globalPinnedUsage / numBoundedInstances
+      } else {
+        0L
+      }
+    } else {
+      // In unbounded memory mode, each instance has its own cache
+      requestedPinnedUsage
+    }
+  }
+
   def getOrCreateRocksDBMemoryManagerAndCache(conf: RocksDBConf): (WriteBufferManager, Cache)
     = synchronized {
     // Register with UnifiedMemoryManager (idempotent operation)

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

@@ -3791,6 +3791,53 @@ class RocksDBSuite extends AlsoTestWithRocksDBFeatures with SharedSparkSession
   }
 
   def listFiles(file: String): Seq[File] = listFiles(new File(file))
+
+  test("SPARK-53792: RocksDBMemoryManager getInstancePinnedBlocksMemUsage") {
+    try {
+      // Clear any existing providers from previous tests
+      RocksDBMemoryManager.resetWriteBufferManagerAndCache
+
+      val boundedMemoryId1 = "test-instance-1"
+      val boundedMemoryId2 = "test-instance-2"
+      val unboundedMemoryId = "test-instance"
+      val sharedCachePinnedUsage = 1000L
+      val requestedPinnedUsage = 300L  // This should be ignored for bounded memory
+
+      // Register two bounded memory instances
+      RocksDBMemoryManager.updateMemoryUsage(boundedMemoryId1, 0L, isBoundedMemory = true)
+      RocksDBMemoryManager.updateMemoryUsage(boundedMemoryId2, 0L, isBoundedMemory = true)
+      RocksDBMemoryManager.updateMemoryUsage(unboundedMemoryId, 0L, isBoundedMemory = false)
+
+      // Test that both instances get the same divided value from globalPinnedUsage
+      val result1 = RocksDBMemoryManager.getInstancePinnedBlocksMemUsage(
+        boundedMemoryId1,
+        sharedCachePinnedUsage,
+        requestedPinnedUsage)
+      val result2 = RocksDBMemoryManager.getInstancePinnedBlocksMemUsage(
+        boundedMemoryId2,
+        sharedCachePinnedUsage,
+        requestedPinnedUsage)
+      val result3 = RocksDBMemoryManager.getInstancePinnedBlocksMemUsage(
+        unboundedMemoryId,
+        sharedCachePinnedUsage,
+        requestedPinnedUsage)
+
+      // With 2 bounded instances, each should get half of globalPinnedUsage
+      assert(result1 === 500L, s"Expected 500L for bounded instance 1, got $result1")
+      assert(result2 === 500L, s"Expected 500L for bounded instance 2, got $result2")
+      assert(result3 === 300L, s"Expected 300L for unbounded instance, got $result3")
+
+      // Test with zero instances (unregistered instance)
+      RocksDBMemoryManager.resetWriteBufferManagerAndCache
+      val resultZero = RocksDBMemoryManager.getInstancePinnedBlocksMemUsage(
+        boundedMemoryId1,
+        sharedCachePinnedUsage,
+        requestedPinnedUsage)
+      assert(resultZero === 0L, s"Expected 0L when no instances, got $resultZero")
+    } finally {
+      RocksDBMemoryManager.resetWriteBufferManagerAndCache
+    }
+  }
 }
 
 object RocksDBSuite {