address comments

Author: Ubuntu

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

@@ -1495,10 +1495,9 @@ 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 totalPinnedBlocksMemUsage = lruCache.getPinnedUsage()
     val pinnedBlocksMemUsage = RocksDBMemoryManager.getInstancePinnedBlocksMemUsage(
-      instanceUniqueId, globalPinnedBlocksMemUsage, requestedPinnedBlocksMemUsage)
+      instanceUniqueId, totalPinnedBlocksMemUsage)
 
     RocksDBMetrics(
       numKeysOnLoadedVersion,

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

@@ -117,9 +117,6 @@ 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
     }
   }
@@ -133,26 +130,17 @@ object RocksDBMemoryManager extends Logging with UnmanagedMemoryConsumer {
    */
   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)
+      pinnedUsage: Long): Long = {
+    val instanceInfo = instanceMemoryMap.
+      getOrDefault(uniqueId, InstanceMemoryInfo(0L, isBoundedMemory = false))
     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
-      }
+      pinnedUsage / numBoundedInstances
     } else {
       // In unbounded memory mode, each instance has its own cache
-      requestedPinnedUsage
+      pinnedUsage
     }
   }
 

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

@@ -2668,6 +2668,53 @@ class RocksDBSuite extends AlsoTestWithRocksDBFeatures with SharedSparkSession
     }
   }
 
+  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 cacheUsage = 1000L
+      val cacheUsage1 = 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,
+        cacheUsage)
+      val result2 = RocksDBMemoryManager.getInstancePinnedBlocksMemUsage(
+        boundedMemoryId2,
+        cacheUsage)
+      val result3 = RocksDBMemoryManager.getInstancePinnedBlocksMemUsage(
+        unboundedMemoryId,
+        cacheUsage1)
+
+      // 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 nonexistInstanceRes = RocksDBMemoryManager.getInstancePinnedBlocksMemUsage(
+        boundedMemoryId1,
+        cacheUsage)
+      assert(
+        nonexistInstanceRes === cacheUsage,
+        s"Expected $cacheUsage when no instances, got $nonexistInstanceRes"
+      )
+    } finally {
+      RocksDBMemoryManager.resetWriteBufferManagerAndCache
+    }
+  }
+
+
   testWithColumnFamilies("SPARK-37224: flipping option 'trackTotalNumberOfRows' during restart",
     TestWithBothChangelogCheckpointingEnabledAndDisabled) { colFamiliesEnabled =>
     withTempDir { dir =>
@@ -3672,52 +3719,6 @@ 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 {