restart iterator on exceptions in Scala

Question

I have an iterator (actually a Source.getLines) that's reading an infinite stream of data from a URL. Occasionally the iterator throws a java.io.IOException when there is a connection problem. In such situations, I need to re-connect and re-start the iterator. I want this to be seamless so that the iterator just looks like a normal iterator to the consumer, but underneath is restarting itself as necessary.

For example, I'd like to see the following behavior:

scala> val iter = restartingIterator(() => new Iterator[Int]{
  var i = -1
  def hasNext = {
    if (this.i < 3) {
      true
    } else {
      throw new IOException
    }
  }
  def next = {
    this.i += 1
    i
  }
})
res0: ...

scala> iter.take(6).toList
res1: List[Int] = List(0, 1, 2, 3, 0, 1)

I have a partial solution to this problem, but it will fail on some corner cases (e.g. an IOException on the first item after a restart) and it's pretty ugly:

def restartingIterator[T](getIter: () => Iterator[T]) = new Iterator[T] {
  var iter = getIter()
  def hasNext = {
    try {
      iter.hasNext
    } catch {
      case e: IOException => {
        this.iter = getIter()
        iter.hasNext
      }
    }
  }
  def next = {
    try {
      iter.next
    } catch {
      case e: IOException => {
        this.iter = getIter()
        iter.next
      }
    }
  }
}

I keep feeling like there's a better solution to this, maybe some combination of Iterator.continually and util.control.Exception or something like that, but I couldn't figure one out. Any ideas?

Answer 1

There is a better solution, the Iteratee:

http://apocalisp.wordpress.com/2010/10/17/scalaz-tutorial-enumeration-based-io-with-iteratees/

Here is for example an enumerator that restarts on encountering an exception.

def enumReader[A](r: => BufferedReader, it: IterV[String, A]): IO[IterV[String, A]] = {
  val tmpReader = r
  def loop: IterV[String, A] => IO[IterV[String, A]] = {
    case i@Done(_, _) => IO { i }
    case Cont(k) => for {
      s <- IO { try { val x = tmpReader.readLine; IO(x) }
                catch { case e => enumReader(r, it) }}.join
      a <- if (s == null) k(EOF) else loop(k(El(s)))
    } yield a
  }
  loop(it)
}

The inner loop advances the Iteratee, but the outer function still holds on to the original. Since Iteratee is a persistent data structure, to restart you just have to call the function again.

I'm passing the Reader by name here so that r is essentially a function that gives you a fresh (restarted) reader. In practise you will want to bracket this more effectively (close the existing reader on exception).

Answer 2

This is fairly close to your version and using scala.util.control.Exception:

def restartingIterator[T](getIter: () => Iterator[T]) = new Iterator[T] {
  import util.control.Exception.allCatch
  private[this] var i = getIter()
  private[this] def replace() = i = getIter()
  def hasNext: Boolean = allCatch.opt(i.hasNext).getOrElse{replace(); hasNext}
  def next(): T = allCatch.opt(i.next).getOrElse{replace(); next}
}

For some reason this is not tail recursive but it that can be fixed by using a slightly more verbose version:

def restartingIterator2[T](getIter: () => Iterator[T]) = new Iterator[T] {
  import util.control.Exception.allCatch
  private[this] var i = getIter()
  private[this] def replace() = i = getIter()
  @annotation.tailrec def hasNext: Boolean = {
    val v = allCatch.opt(i.hasNext)
    if (v.isDefined) v.get else {replace(); hasNext}
  }
  @annotation.tailrec def next(): T = {
    val v = allCatch.opt(i.next)
    if (v.isDefined) v.get else {replace(); next}
  }
}

Answer 3

Here's an answer that doesn't work, but feels like it should:

def restartingIterator[T](getIter: () => Iterator[T]): Iterator[T] = {
  new Traversable[T] {
    def foreach[U](f: T => U): Unit = {
      try {
        for (item <- getIter()) {
          f(item)
        }
      } catch {
        case e: IOException => this.foreach(f)
      }
    }
  }.toIterator
}

I think this very clearly describes the control flow, which is great.

This code will throw a StackOverflowError in Scala 2.8.0 because of a bug in Traversable.toStream, but even after the fix for that bug, this code still won't work for my use case because toIterator calls toStream, which means that it will store all items in memory.

I'd love to be able to define an Iterator by just writing a foreach method, but there doesn't seem to be any easy way to do that.