[SPARK-38640][CORE] Fix NPE with memory-only cache blocks and RDD fetching

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

Fixes a bug where if `spark.shuffle.service.fetch.rdd.enabled=true`, memory-only cached blocks will fail to unpersist.

### Why are the changes needed?

In #33020, when all RDD blocks are removed from `externalShuffleServiceBlockStatus`, the underlying Map is nulled to reduce memory. When persisting blocks we check if it's using disk before adding it to `externalShuffleServiceBlockStatus`, but when removing them there is no check, so a memory-only cache block will keep `externalShuffleServiceBlockStatus` null, and when unpersisting it throw an NPE because it tries to remove from the null Map. This adds checks to the removal as well to only remove if the block is on disk, and therefore should have been added to `externalShuffleServiceBlockStatus` in the first place.

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

No

### How was this patch tested?

New and updated UT

Closes #35959 from Kimahriman/fetch-rdd-memory-only-unpersist.

Authored-by: Adam Binford <[email protected]>
Signed-off-by: Sean Owen <[email protected]>

Author: Kimahriman

core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala

@@ -838,9 +838,11 @@ private[spark] class BlockStatusPerBlockId {
   }
 
   def remove(blockId: BlockId): Unit = {
-    blocks.remove(blockId)
-    if (blocks.isEmpty) {
-      blocks = null
+    if (blocks != null) {
+      blocks.remove(blockId)
+      if (blocks.isEmpty) {
+        blocks = null
+      }
     }
   }
 

core/src/test/scala/org/apache/spark/ExternalShuffleServiceSuite.scala

@@ -255,4 +255,26 @@ class ExternalShuffleServiceSuite extends ShuffleSuite with BeforeAndAfterAll wi
       }
     }
   }
+
+  test("SPARK-38640: memory only blocks can unpersist using shuffle service cache fetching") {
+    for (enabled <- Seq(true, false)) {
+      val confWithRddFetch =
+        conf.clone.set(config.SHUFFLE_SERVICE_FETCH_RDD_ENABLED, enabled)
+      sc = new SparkContext("local-cluster[1,1,1024]", "test", confWithRddFetch)
+      sc.env.blockManager.externalShuffleServiceEnabled should equal(true)
+      sc.env.blockManager.blockStoreClient.getClass should equal(classOf[ExternalBlockStoreClient])
+      try {
+        val rdd = sc.parallelize(0 until 100, 2)
+          .map { i => (i, 1) }
+          .persist(StorageLevel.MEMORY_ONLY)
+
+        rdd.count()
+        rdd.unpersist(true)
+        assert(sc.persistentRdds.isEmpty)
+      } finally {
+        rpcHandler.applicationRemoved(sc.conf.getAppId, true)
+        sc.stop()
+      }
+    }
+  }
 }

core/src/test/scala/org/apache/spark/storage/BlockManagerInfoSuite.scala

@@ -63,6 +63,8 @@ class BlockManagerInfoSuite extends SparkFunSuite {
     if (svcEnabled) {
       assert(getEssBlockStatus(bmInfo, rddId).isEmpty)
     }
+    bmInfo.updateBlockInfo(rddId, StorageLevel.NONE, memSize = 0, diskSize = 0)
+    assert(bmInfo.remainingMem === 30000)
   }
 
   testWithShuffleServiceOnOff("RDD block with MEMORY_AND_DISK") { (svcEnabled, bmInfo) =>