akka.actor.Timers Scala Examples

package ai.forestflow.serving.cluster

import java.io.File

import akka.actor.{Actor, ActorLogging, ActorRef, Props, Timers}
import akka.cluster.Cluster
import akka.cluster.pubsub.DistributedPubSub
import akka.cluster.pubsub.DistributedPubSubMediator.Subscribe
import ai.forestflow.domain.CleanupLocalStorage
import org.apache.commons.io.FileUtils
import com.typesafe.scalalogging.LazyLogging
import ai.forestflow.utils.ThrowableImplicits._

import scala.util.{Failure, Success, Try}

/***
 * This actor is responsible for node-level (host-level) stuff that should be done on a per-node basis.
 * A good example of this is file system cleanup tasks.
 */
object NodeActor extends LazyLogging {
  
  def props(): Props =
    Props(new NodeActor)
      .withDispatcher("blocking-io-dispatcher")

  def cleanupLocalStorage(path: String): Unit = {
    val localDir = new File(path)
    val localDirExists = localDir.exists()
    logger.info(s"Cleaning up local storage: Local Directory: $localDir , exists? $localDirExists")
    if (localDirExists)
      Try(FileUtils.deleteDirectory(localDir)) match {
        case Success(_) => logger.info(s"Local Directory $localDir cleaned up successfully")
        case Failure(ex) => logger.error(s"Local Directory $localDir cleanup failed! Reason: ${ex.printableStackTrace}")
      }
  }
}

class NodeActor extends Actor
  with ActorLogging
  with Timers {

  
  implicit val cluster: Cluster = Cluster(context.system)
  val mediator: ActorRef = DistributedPubSub(context.system).mediator

  mediator ! Subscribe(classOf[CleanupLocalStorage].getSimpleName, self)

  override def receive: Receive = {
    case CleanupLocalStorage(path) =>
      NodeActor.cleanupLocalStorage(path)
  }
} 

package io.hydrosphere.mist.utils.akka

import akka.pattern.pipe
import akka.actor.{Actor, ActorLogging, ActorRef, ActorRefFactory, Props, ReceiveTimeout, SupervisorStrategy, Terminated, Timers}
import io.hydrosphere.mist.utils.Logger

import scala.concurrent.{Future, Promise}
import scala.concurrent.duration._

class RestartSupervisor(
  name: String,
  start: () => Future[ActorRef],
  timeout: FiniteDuration,
  maxRetry: Int
) extends Actor with ActorLogging with Timers {

  override def receive: Receive = init

  import context._
  import RestartSupervisor._

  private def init: Receive = {
    case Event.Start(req) =>
      start().map(Event.Started) pipeTo self
      context become await(Some(req), 0)
  }

  private def await(req: Option[Promise[ActorRef]], attempts: Int): Receive = {
    case Event.Started(ref) =>
      req.foreach(_.success(self))
      context watch ref
      context become proxy(ref)

    case akka.actor.Status.Failure(e)  if maxRetry == attempts + 1 =>
      req.foreach(_.failure(e))
      log.error(e, "Starting child for {} failed, maxRetry reached", name)
      context stop self

    case akka.actor.Status.Failure(e) =>
      log.error(e, "Starting child for {} failed", name)
      timers.startSingleTimer("timeout", Event.Timeout, timeout)
      context become restartTimeout(req, attempts)
  }

  private def proxy(ref: ActorRef): Receive = {
    case Terminated(_) =>
      log.error(s"Reference for {} was terminated. Restarting", name)
      timers.startSingleTimer("timeout", Event.Timeout, timeout)
      context become restartTimeout(None, 0)

    case x => ref.forward(x)
  }

  private def restartTimeout(req: Option[Promise[ActorRef]], attempts: Int): Receive = {
    case Event.Timeout =>
      start().map(Event.Started) pipeTo self
      context become await(req, attempts + 1)
  }
}

object RestartSupervisor {

  sealed trait Event
  object Event {
    final case class Start(req: Promise[ActorRef]) extends Event
    case object Restart extends Event
    final case class Started(ref: ActorRef) extends Event
    case object Timeout extends Event
  }


  def props(
    name: String,
    start: () => Future[ActorRef],
    timeout: FiniteDuration,
    maxRetry: Int
  ): Props = {
    Props(classOf[RestartSupervisor], name, start, timeout, maxRetry)
  }

  def wrap(
    name: String,
    start: () => Future[ActorRef],
    timeout: FiniteDuration,
    maxRetry: Int
  )(implicit af: ActorRefFactory): Future[ActorRef] = {

    val ref = af.actorOf(props(name, start, timeout, maxRetry))
    val promise = Promise[ActorRef]
    ref ! Event.Start(promise)
    promise.future
  }

  def wrap(
    name: String,
    maxRetry: Int,
    start: () => Future[ActorRef]
  )(implicit af: ActorRefFactory): Future[ActorRef] = wrap(name, start, 5 seconds, maxRetry)(af)

}