Spark - Cleaning Up Objects with WeakReference

I recently saw a question on Stack Overflow about when Spark reclaims variables such as Accumulators and Broadcasts. While reading the source, I came across this interesting mechanism.

Spark ContextCleaner

When Spark is idle, with no tasks running, we often see logs like these:

1
19/02/12 05:19:51 INFO Spark Context Cleaner org.apache.spark.internal.Logging$class.logInfo(Logging.scala:54): Cleaned accumulator 108284023
2
19/02/12 05:19:51 INFO Spark Context Cleaner org.apache.spark.internal.Logging$class.logInfo(Logging.scala:54): Cleaned accumulator 108283658

Tracing the source reveals that SparkContext initializes ContextCleaner on startup and launches a daemon cleaningThread. This thread loops continuously, cleaning up unused objects such as RDDs, Broadcast variables, and Accumulators.
This raises a question: taking an Accumulator as an example, how does ContextCleaner know that it is no longer used, meaning that no object references it?

First, let’s look at the cleanup process:

1
val reference = Option(referenceQueue.remove(ContextCleaner.REF_QUEUE_POLL_TIMEOUT))
2
  .map(_.asInstanceOf[CleanupTaskWeakReference])
3
// Synchronize here to avoid being interrupted on stop()
4
synchronized {
5
  reference.foreach { ref =>
6
    logDebug("Got cleaning task " + ref.task)
7
    referenceBuffer.remove(ref)
8
    ref.task match {
9
      case CleanRDD(rddId) =>
10
        doCleanupRDD(rddId, blocking = blockOnCleanupTasks)
11
      case CleanShuffle(shuffleId) =>
12
        doCleanupShuffle(shuffleId, blocking = blockOnShuffleCleanupTasks)
13
      case CleanBroadcast(broadcastId) =>
14
        doCleanupBroadcast(broadcastId, blocking = blockOnCleanupTasks)
15
      case CleanAccum(accId) =>
16
        doCleanupAccum(accId, blocking = blockOnCleanupTasks)
17
      case CleanCheckpoint(rddId) =>
18
        doCleanCheckpoint(rddId)
19
    }
20
  }
21
}

ContextCleaner finds the object to clean up (CleanAccum) through a referenceQueue. Starting there, let’s examine WeakReference in the JVM.

Java WeakReference

Java has several kinds of references:

• StrongReference: The objects we normally define use these references and are harder to collect.
• WeakReference: For example, Spark wraps references in CleanupTaskWeakReference(task, objectForCleanup, referenceQueue). If the referenced object (task) is reachable only through this WeakReference, it is reclaimed during GC and placed in referenceQueue.
• SoftReference: Stronger than a WeakReference. The referent can be reclaimed, but not necessarily during the next GC.

ContextCleaner therefore combines WeakReference with referenceQueue to reclaim objects. Registering an Accumulator also calls registerForCleanup:

1
/** Register an object for cleanup. */
2
private def registerForCleanup(objectForCleanup: AnyRef, task: CleanupTask): Unit = {
3
  referenceBuffer.add(new CleanupTaskWeakReference(task, objectForCleanup, referenceQueue))
4
}

referenceBuffer prevents the WeakReference itself from being collected before it is processed.

Spark wraps the registered Accumulator in a CleanupTask and initializes a WeakReference based on that task. Once the Accumulator is no longer referenced, the task is placed in referenceQueue. The cleaningThread retrieves the object awaiting collection from the queue and processes it, as shown in the cleanup code above.