SlowIterationLogger

example_shell.scala

// an example of using SlowLoggingIterator, with the silly SlowIterator example, from a scala shell.
// you could paste this in after invoking the shell with
// > scala -i SlowLoggingIterator.scala -i SlowLoggingIteratorExample.scala
import scala.collection.JavaConverters._
val originalSlowItr = new SlowIterator(100, 500, 20, 10)

var x = 0
val loggingItr = new SlowLoggingIterator(originalSlowItr, 100, 1000, t => s"state info: x = $x")

loggingItr.asScala.foreach { i =>
  x += 500
  println(i)
}

SlowIterator.scala

// This is an example iterator that runs slowly, to demonstrate how SlowLoggingIterator works
// it just iterates over a range of ints, but puts in occassional delays, to simulate an iterator that is
// actually doing something more complex, eg. fetching records from a DB which is occassionaly slow.

class SlowIterator(start: Int, end: Int, delay: Long, every: Int) extends java.util.Iterator[Integer] {
  val underlying = (start until end).toIterator

  def hasNext(): Boolean = underlying.hasNext

  def next(): Integer = {
    val n = underlying.next()
    if (n % every == 0) {
      Thread.sleep(delay)
    }
    n
  }
}

SlowLoggingIterator.scala

import java.util.Timer
import java.util.TimerTask
import java.util.concurrent.atomic._

object SlowLoggingIterator {
  val timer = new Timer("slow-logging-timer", true)
}

class SlowLoggingIterator[T <: AnyRef](wrapped: java.util.Iterator[T], slowLogMillis: Long, stateLogMillis: Long, stateInfo: Long => String) extends java.util.Iterator[T] {
  val lastRecord = new AtomicReference[T]()
  val recordCount = new AtomicInteger(0)
  val totalTime = new AtomicLong(0)
  val stateLogTask = new TimerTask() {
      def run(): Unit = logCurrentState()
  }
  SlowLoggingIterator.timer.scheduleAtFixedRate(stateLogTask, stateLogMillis, stateLogMillis)

  private def runWithTimer[R](f: () => R): R = {
    val start = System.currentTimeMillis()
    val timer = new SlowLoggingTimerTask(lastRecord.get(), start)
    SlowLoggingIterator.timer.scheduleAtFixedRate(timer, slowLogMillis, slowLogMillis)
    val r = f()
    timer.cancel()
    val end = System.currentTimeMillis()
    totalTime.getAndAdd(end - start)
    r
  }

  def hasNext(): Boolean = {
    val r = runWithTimer { () => wrapped.hasNext() }
    if (!r) {
      stateLogTask.cancel()
    }
    r
  }

  def next(): T = {
    val r = runWithTimer { () => wrapped.next() }
    lastRecord.set(r)
    recordCount.getAndIncrement()
    r
  }

  def logCurrentState(): Unit = {
    val t = totalTime.get()
    log(s"${recordCount.get()} records read; $t ms in iterator in total; ${stateInfo(t)}")
  }

  def logNoProgress(last: T, startTime: Long, now: Long): Unit = {
    val time = now - startTime
    log(s"No progress in last $time.  Last record fetched is $last")
    logCurrentState()
  }

  // override this if you want something other than println, eg. to use log4j
  def log(msg: => String): Unit = {
    println(msg)
  }

  private class SlowLoggingTimerTask(lastVal: T, startTime: Long) extends TimerTask {
    def run(): Unit = {
      if (lastVal eq lastRecord.get()) {
        logNoProgress(lastVal, startTime, System.currentTimeMillis())
      }
    }
  }
}