HeartMonitor using Akka FSM

gistfile1.scala

import akka.actor.{Actor, FSM}
import java.util.concurrent.TimeUnit

sealed trait Health
case object Stale extends Health
case object Alive extends Health
case object Dead extends Health

case object HeartBeat

class HeartMonitor(deadAfter: Long, checkEvery: Long, timeUnit: TimeUnit) extends Actor with FSM[Health, Long] {
  import System.{currentTimeMillis => now}
  val checkTimeout = (checkEvery, timeUnit)

  notifying {
    case Transition(Stale, Alive) =>
      log.info("HeartMonitor received initial heartbeat")
    case Transition(Dead, Alive) =>
      log.info("HeartMonitor noticed we are back alive again")
    case Transition(_, Dead) =>
      log.error("HeartMonitor thinks we are dead")
  }

  when(Stale) {
    case Beat() => goto(Alive) using now forMax checkTimeout
    case SkippedBeat() => stay forMax checkTimeout
    case NoBeat() => goto(Dead)
  }

  when(Alive) {
    case Beat() => stay using now forMax checkTimeout
    case SkippedBeat() => stay forMax checkTimeout
    case NoBeat() => goto(Dead)
  }

  when(Dead) {
    case Beat() => goto(Alive) using now forMax checkTimeout
  }

  startWith(Stale, now, Some(checkTimeout))

  object SkippedBeat {
    def unapply(event: Event) = event match {
      case Event(StateTimeout, lastHeartBeatTimeMillis) if !isDead(lastHeartBeatTimeMillis) => true
      case _ => false
    }
  }

  object NoBeat {
    def unapply(event: Event) = event match {
      case Event(StateTimeout, lastHeartBeatTimeMillis) if isDead(lastHeartBeatTimeMillis) => true
      case _ => false
    }
  }

  object Beat {
    def unapply(event: Event) = event match {
      case Event(HeartBeat, _) => true
      case _ => false
    }
  }

  def isDead(lastHeartBeatTimeMillis: Long): Boolean = {
    now - lastHeartBeatTimeMillis > timeUnit.toMillis(deadAfter)
  }
}