Distributing workload using Akka Router with exception handling

Question

I am fairly new to Akka and I am interested in any feedback about the following code. Is this the correct way to go about distributing workload across a number of threads? Do I use the Akka API correctly? Is this the right pattern for handling the routees exceptions?

package myakka.routing

import java.util.concurrent.TimeUnit

import scala.collection.mutable.{ Map => MMap }
import scala.concurrent.Await
import scala.concurrent.duration.Duration
import scala.concurrent.duration.DurationInt
import scala.util.Random

import akka.actor.Actor
import akka.actor.ActorContext
import akka.actor.ActorLogging
import akka.actor.ActorRef
import akka.actor.ActorSystem
import akka.actor.PoisonPill
import akka.actor.Props
import akka.actor.ReceiveTimeout
import akka.actor.Terminated
import akka.actor.actorRef2Scala
import akka.pattern.ask
import akka.routing.ActorRefRoutee
import akka.routing.RoundRobinRoutingLogic
import akka.routing.Router
import akka.util.Timeout

trait Message
case class Work(w: String, listenerRef: ActorRef) extends Message
case object WorkDone extends Message
case class StartWork(workload: Seq[String]) extends Message
case object WorkAvailable extends Message
case object AvailableForWork extends Message
case object RequestResult extends Message

/*
 * This sample uses an Akka Router
 */

/*
 * This mixin can be used for adding timeout to an actor
 */
trait TimedActor {
  this: Actor =>

  val defaultTimeout = 5 seconds

  context.setReceiveTimeout(defaultTimeout)

  val timeoutBehaviour: Receive = {
    case ReceiveTimeout =>
      context.setReceiveTimeout(Duration.Undefined)
      throw new RuntimeException("Received timeout")
  }
}

Dispatcher:

object Dispatcher {
  def props(ref: ActorRef) = Props(classOf[Dispatcher], ref)

  private def makeRouter(ctx: ActorContext): Router = {
    val routees = Vector.fill(Runtime.getRuntime().availableProcessors()) {
      val r = ctx.actorOf(Worker.props)

      ActorRefRoutee(r)
    }

    Router(RoundRobinRoutingLogic(), routees)
  }
}

class Dispatcher(listener: ActorRef) extends Actor with ActorLogging {
  var iterator: Iterator[String] = Iterator.empty

  var router = Dispatcher.makeRouter(context)

  /*
   * This is our routine behaviour
   */
  def dispatcherBehavior: Receive = {
    case AvailableForWork =>
      log.info(s"Worker available for work: $sender")
      checkForWork(sender)

    case StartWork(workload) => {
      iterator = workload.iterator
      log.info(s"Starting work as requested by $sender")
      router.routees.foreach(_ => router.route(WorkAvailable, listener)) // let all routees know that work is available    
    }

    case e @ _ => log.error(s"Unknown message $e received from $sender in dispatcherBehavior")
  }

  def receive = dispatcherBehavior

  private def checkForWork(workerRef: ActorRef) = {
    if (iterator.isEmpty) {
      log.info(s"No work to do!")
      listener ! WorkDone
    } else {
      val item = iterator.next
      log.info(s"Sending another piece of work $item to $workerRef")
      workerRef ! Work(item, listener)
    }
  }

}

Worker:

object Worker {
  def props = Props(classOf[Worker])
}

class Worker extends Actor with ActorLogging {
  import Listener.UnitCompleted

  override def postRestart(reason: Throwable) = {
    preStart
    log.info("Reporting for duty after restart")
    context.parent ! AvailableForWork
  }

  def workerBehaviour: Receive = {
    case WorkAvailable =>
      log.info(s"Notified about available work by ${context.parent}")
      context.parent ! AvailableForWork

    case Work(w, listener) =>
      log.info(s"Got work $w from ${context.parent}")
      listener ! UnitCompleted(w, doWork(w))
      log.info(s"Completed work $w")
      if (!context.system.isTerminated) {
        log.info(s"Letting master know we, $self, are avaible for work.")
        context.parent ! AvailableForWork
      }

    case e @ _ => log.error(s"Unknown message $e received in workerBehaviour from $sender")
  }

  def receive = workerBehaviour

  private def doWork(item: String): Int = {
    log.info(s"Doing the work=$item")

    // Dummy work
    val result = {
      (1 to 2000000).reduceLeft {
        _ + _
      }

      item.toInt
    }

    val fail = Random.nextInt(1000) match {
      case x if ((0 to 500) contains x) => false
      case _ => true
    }

    if (fail) {
      throw new Exception("Don't feel like working today!")
    }

    log.info(s"Finished the work=$item: $result")

    result
  }

}

Listener:

object Listener {
  case class UnitCompleted(id: String, result: Int)

  def props(workload: Seq[String], timeout: Timeout) = Props(classOf[Listener], workload, timeout)

}

class Listener(workload: Seq[String], timeout: Timeout) extends Actor with ActorLogging {
  import Listener.UnitCompleted
  import scala.collection.mutable.{ Map => MMap }

  var resultTarget: ActorRef = _

  context.setReceiveTimeout(timeout.duration)

  val workDone: MMap[String, Option[Int]] = {
    var m: MMap[String, Option[Int]] = MMap()
    workload.foreach(k => m(k) = None)
    m
  }

  def receive = {
    case UnitCompleted(id, r) =>
      log.info(s"Work completed for $id with result=$r by $sender")
      workDone(id) = Some(r)
      if (checkWorkDone) {
        log.info("All work done!")
        resultTarget ! workDone.values.flatten
      }

    case RequestResult =>
      resultTarget = sender

    case WorkDone =>
      if (!checkWorkDone) {
        log.info("Still more work todo!")
        sender ! StartWork(workTodo)
      }

    case ReceiveTimeout =>
      log.error("Received timeout, shutting down!")
      context.system.shutdown
  }

  private def checkWorkDone: Boolean = {
    !workDone.valuesIterator.exists(_ == None)
  }

  private def workTodo: Seq[String] = workDone.filter({ case (k, v) => v == None }).keys.toSeq
}

WorkService:

class WorkService(implicit t: Timeout) {
  def doWork(workload: Seq[String]): Seq[Option[Int]] = {
    val system = ActorSystem("workpull")
    val listener = system.actorOf(Listener.props(workload, t), "listener")
    val dispatcher = system.actorOf(Dispatcher.props(listener), "dispatcher")

    dispatcher.tell(StartWork(workload), listener)

    /*
     *  (listener ? RequestResult) is equivalent to ask(listener, RequestResult) and returns a Future.
     *  Await.result waits for the Future to return the actual value.
     */
    val result = Await.result(listener ? RequestResult, t.duration).asInstanceOf[Seq[Option[Int]]]

    system.shutdown

    result
  }
}

WorkDistributor:

object WorkDistributor extends App {
  val workload = Seq("1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12")

  implicit val timeout = Timeout(15, TimeUnit.SECONDS)
  val result = new WorkService doWork (workload)
  println("Final result:" + result)
}

Answer 1

This is very well written example code, well done! It is also complete, I was able to compile and execute the code with no changes. Very nice.

Is this the correct way to go about distributing workload across a number of threads?

Yes. The only improvement I would suggest in this regard is that I would rather see the number of routees specified in a configuration file than in the code. This would allow you to reconfigure the router in a number of different ways without having to recompile the application. You could also deploy the service across multiple machines just through configuration. Limiting the routees to the number of available processors makes the assumption that the routees are all deployed locally, and this runs counter to the principle of location transparency.

So a configuration file would be good for this example, especially for the router.

Do I use the Akka API correctly?

Yes.

Is this the right pattern for handling the routees exceptions?

As far as this example goes, yes, it seems to be perfectly correct. It uses the default behavior of actors being restarted when an exception is thrown, and in this example that seems to be exactly what is needed. In other situations you may have to change the default supervisor strategy and, for example, escalate the exception to the parent for special handling instead of just letting the actor be restarted.

---

Aside from a functional review, and you said you were interested in "any feedback", I will also add some suggestions on how you might make the code (even) cleaner. Some of this might be personal preference, so feel free to disagree:*

• It appears the entire listing is in one file. Maybe you did that just to submit it for review, but it would be easier to digest in separate files, each with their own imports. The way it was split up into Worker, Dispatcher, etc., would be a good way to split it into separate files as well.
• There are number of unused imports, this is cruft.
• The TimedActor trait is not actually used anywhere, so it doesn't add value to the example.
• I would use parentheses after shutdown() and preStart() since these methods potentially have side effects. See Programming in Scala 10.3 Defining parameterless methods.
• The curly brackets are not necessary after case. I personally prefer to leave them out where they are not necessary.

• The reduceLeft call can be simplified:

  From

  (1 to 2000000).reduceLeft {
     _ + _
  }
  

  To

  (1 to 2000000).sum
  

• Unnecessary parentheses in a case filter:

    case x if ((0 to 500) contains x) ⇒ false
  

  can be written as

    case x if (0 to 500) contains x ⇒ false
  

• Simplify

  !workDone.valuesIterator.exists(_ == None)
  

  to

  !workDone.valuesIterator.contains(None)
  

• The following code

    val result = new WorkService doWork (workload)
  

  might be nicer written as

    val result = new WorkService.doWork(workload)
  

  or

    val result = new WorkService doWork workload
  

  (Personally I prefer the former).

---

* With credit given to the inspections in the Scala Plugin for IntelliJ for automatically discovering many of these issues