Expose Spark's cache memory status as Metrics Gauges

SparkContextInstrumentor.scala

import com.yammer.metrics.Metrics
import com.yammer.metrics.core.Gauge
import org.apache.spark.SparkContext

/**
 * Created by tzachz on 10/21/15
 */
object SparkContextInstrumentor {

  def instrument(context: SparkContext): SparkContext = {
    registerGauge("maxCacheMemory",        u => context.getExecutorMemoryStatus.map(_._2._1).sum)
    registerGauge("freeCacheMemory",       u => context.getExecutorMemoryStatus.map(_._2._2).sum)
    registerGauge("cachedRDDs",            u => context.getRDDStorageInfo.count(_.isCached))
    registerGauge("cachedRDDsTotalMemory", u => context.getRDDStorageInfo.filter(_.isCached).map(_.memSize).sum)
    context
  }

  private def registerGauge(name: String, valueProducer: Unit => Long): Unit = {
    Metrics.newGauge[Long](classOf[SparkContext], name, new Gauge[Long]() {
      override def value(): Long = {
        valueProducer.apply()
      }
    })
  }

}

This deserves a blog post 😄 and I got here via Google, believe it or not.