[SPARK-26726] Synchronize the amount of memory used by the broadcast variable to the UI display

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

@@ -759,18 +759,10 @@ private[spark] class AppStatusListener(
 
       // Use RDD distribution to update executor memory and disk usage info.
       liveRDD.getDistributions().foreach { case (executorId, rddDist) =>
-        liveExecutors.get(executorId).foreach { exec =>
-          if (exec.hasMemoryInfo) {
-            if (storageLevel.useOffHeap) {
-              exec.usedOffHeap = addDeltaToValue(exec.usedOffHeap, -rddDist.offHeapUsed)
-            } else {
-              exec.usedOnHeap = addDeltaToValue(exec.usedOnHeap, -rddDist.onHeapUsed)
-            }
-          }
-          exec.memoryUsed = addDeltaToValue(exec.memoryUsed, -rddDist.memoryUsed)
-          exec.diskUsed = addDeltaToValue(exec.diskUsed, -rddDist.diskUsed)
-          maybeUpdate(exec, now)
-        }
+        val maybeExec = liveExecutors.get(executorId)
+        updateExecutorMemoryDiskInfo(maybeExec, storageLevel, -rddDist.memoryUsed,
+          -rddDist.diskUsed, Option(-rddDist.offHeapUsed), Option(-rddDist.onHeapUsed))
+        maybeExec.foreach(exec => maybeUpdate(exec, now))
       }
     }
 
@@ -827,6 +819,7 @@ private[spark] class AppStatusListener(
     event.blockUpdatedInfo.blockId match {
       case block: RDDBlockId => updateRDDBlock(event, block)
       case stream: StreamBlockId => updateStreamBlock(event, stream)
+      case broadcast: BroadcastBlockId => updateBroadcastBlock(event, broadcast)
       case _ =>
     }
   }
@@ -884,17 +877,7 @@ private[spark] class AppStatusListener(
 
     // Update the executor stats first, since they are used to calculate the free memory
     // on tracked RDD distributions.
-    maybeExec.foreach { exec =>
-      if (exec.hasMemoryInfo) {
-        if (storageLevel.useOffHeap) {
-          exec.usedOffHeap = addDeltaToValue(exec.usedOffHeap, memoryDelta)
-        } else {
-          exec.usedOnHeap = addDeltaToValue(exec.usedOnHeap, memoryDelta)
-        }
-      }
-      exec.memoryUsed = addDeltaToValue(exec.memoryUsed, memoryDelta)
-      exec.diskUsed = addDeltaToValue(exec.diskUsed, diskDelta)
-    }
+    updateExecutorMemoryDiskInfo(maybeExec, storageLevel, memoryDelta, diskDelta, None, None)
 
     // Update the block entry in the RDD info, keeping track of the deltas above so that we
     // can update the executor information too.
@@ -995,6 +978,44 @@ private[spark] class AppStatusListener(
     }
   }
 
+  private def updateBroadcastBlock(
+      event: SparkListenerBlockUpdated,
+      broadcast: BroadcastBlockId): Unit = {
+    val now = System.nanoTime()
+    val executorId = event.blockUpdatedInfo.blockManagerId.executorId
+    val storageLevel = event.blockUpdatedInfo.storageLevel
+
+    // Whether values are being added to or removed from the existing accounting.
+    val diskDelta = event.blockUpdatedInfo.diskSize * (if (storageLevel.useDisk) 1 else -1)
+    val memoryDelta = event.blockUpdatedInfo.memSize * (if (storageLevel.useMemory) 1 else -1)
+
+    val maybeExec = liveExecutors.get(executorId)
+    updateExecutorMemoryDiskInfo(maybeExec, storageLevel, memoryDelta, diskDelta, None, None)
+    maybeExec.foreach(exec => maybeUpdate(exec, now))
+
+  }
+
+  // update executor memory and disk usage info
+  private def updateExecutorMemoryDiskInfo(
+      maybeExec: Option[LiveExecutor],
+      storageLevel: StorageLevel,
+      memoryDelta: Long,
+      diskDelta: Long,
+      OffHeapDelta: Option[Long],
+      OnHeapDelta: Option[Long]): Unit = {
+    maybeExec.foreach { exec =>
+      if (exec.hasMemoryInfo) {
+        if (storageLevel.useOffHeap) {
+          exec.usedOffHeap = addDeltaToValue(exec.usedOffHeap, OffHeapDelta.getOrElse(memoryDelta))
+        } else {
+          exec.usedOnHeap = addDeltaToValue(exec.usedOnHeap, OnHeapDelta.getOrElse(memoryDelta))
+        }
+      }
+      exec.memoryUsed = addDeltaToValue(exec.memoryUsed, memoryDelta)
+      exec.diskUsed = addDeltaToValue(exec.diskUsed, diskDelta)
+    }
+  }
+
   private def getOrCreateStage(info: StageInfo): LiveStage = {
     val stage = liveStages.computeIfAbsent((info.stageId, info.attemptNumber),
       new Function[(Int, Int), LiveStage]() {

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

@@ -938,6 +938,24 @@ class AppStatusListenerSuite extends SparkFunSuite with BeforeAndAfter {
     intercept[NoSuchElementException] {
       check[StreamBlockData](stream1.name) { _ => () }
     }
+
+    // Update a BroadcastBlock.
+    val broadcast1 = BroadcastBlockId(1L)
+    listener.onBlockUpdated(SparkListenerBlockUpdated(
+      BlockUpdatedInfo(bm1, broadcast1, level, 1L, 1L)))
+
+    check[ExecutorSummaryWrapper](bm1.executorId) { exec =>
+      assert(exec.info.memoryUsed === 1L)
+      assert(exec.info.diskUsed === 1L)
+    }
+
+    // Drop a BroadcastBlock.
+    listener.onBlockUpdated(SparkListenerBlockUpdated(
+      BlockUpdatedInfo(bm1, broadcast1, StorageLevel.NONE, 1L, 1L)))
+    check[ExecutorSummaryWrapper](bm1.executorId) { exec =>
+      assert(exec.info.memoryUsed === 0)
+      assert(exec.info.diskUsed === 0)
+    }
   }
 
   test("eviction of old data") {