JWT Auth Actor state machine using become

JWTAuthActor.scala

import akka.actor.{Actor, ActorRef}
import com.typesafe.scalalogging.LazyLogging
import monix.execution.Scheduler
import org.joda.time.{DateTime, DateTimeZone}

import scala.concurrent.Future
import scala.concurrent.duration._
import scala.util.control.NonFatal
import scala.util.{Failure, Success}


final class JWTAuthActor(
  authenticateClb: () => Future[Either[String, String]],
  refreshTokenClb: () => Future[Either[String, String]],
  retryDelay: FiniteDuration
) (implicit s: Scheduler) extends Actor with LazyLogging {

  import JWTAuthActor._

  def receive: Receive = active(State(None, None))

  def active(state: State): Receive = {
    case Messages.GetToken =>
      state.token match {
        case Some(token) =>
          sender ! Messages.GetTokenSuccess(token)
        case None =>
          if (shouldRetry(state.failedAt)) {
            authenticate()
            context.become(authenticating(sender, state))
          } else
            sender ! Messages.FailedRecentlyRetryLater
      }
    case Messages.RefreshToken =>
      if (shouldRetry(state.failedAt)) {
        refreshToken()
        context.become(refreshingToken(sender, state))
      } else
        sender ! Messages.FailedRecentlyRetryLater
    case Messages.Authenticate =>
      if (shouldRetry(state.failedAt)) {
        authenticate()
        context.become(authenticating(sender, state))
      } else
        sender ! Messages.FailedRecentlyRetryLater
    case unexpected =>
      logger.error(s"Unexpected message received while active: $unexpected")
  }

  def authenticating(originalSender: ActorRef, state: State): Receive = {
    case success @ Messages.AuthenticationSuccess(freshToken) =>
      originalSender ! success
      context.become(active(State(Some(freshToken), None)))
    case error @ Messages.AuthenticationError(_) =>
      originalSender ! error
      context.become(active(State(
        None,
        Some(DateTime.now(DateTimeZone.UTC)))))
    case _ =>
      originalSender ! Messages.AuthenticationInProgress
  }

  def refreshingToken(originalSender: ActorRef, state: State): Receive = {
    case success @ Messages.RefreshTokenSuccess(freshToken) =>
      originalSender ! success
      context.become(active(State(Some(freshToken), None)))
    case Messages.RefreshTokenError(_) =>
      authenticate()
      context.become(authenticating(originalSender, State(
        None,
        Some(DateTime.now(DateTimeZone.UTC)))))
    case _ =>
      originalSender ! Messages.RefreshTokenInProgress
  }

  private[this] def authenticate(): Unit =
    try {
      authenticateClb().onComplete {
        case Success(Right(freshToken)) =>
          self ! Messages.AuthenticationSuccess(freshToken)
        case Success(Left(error)) =>
          self ! Messages.AuthenticationError(error)
        case Failure(throwable) =>
          logger.error(s"Authentication failure", throwable)
          self ! Messages.AuthenticationError(throwable.getMessage)
      }
    } catch {
      case NonFatal(throwable) =>
        logger.error(s"Authentication exception", throwable)
        self ! Messages.AuthenticationError(throwable.getMessage)
    }

  private[this] def refreshToken(): Unit =
    try {
      refreshTokenClb().onComplete {
        case Success(Right(freshToken)) =>
          self ! Messages.RefreshTokenSuccess(freshToken)
        case Success(Left(error)) =>
          logger.error(s"Refresh token error: $error")
          self ! Messages.RefreshTokenError(error)
        case Failure(throwable) =>
          logger.error(s"Refresh token failure", throwable)
          self ! Messages.RefreshTokenError(throwable.getMessage)
      }
    } catch {
      case NonFatal(throwable) =>
        logger.error(s"Refresh token exception", throwable)
        self ! Messages.RefreshTokenError(throwable.getMessage)
    }

  private def shouldRetry(failedAt: Option[DateTime]) =
    failedAt.forall(now.getMillis - _.getMillis > retryDelay.toMillis)
}

object JWTAuthActor {
  object Messages {
    case object GetToken
    case class GetTokenSuccess(token: String)

    case object RefreshToken
    case object RefreshTokenInProgress
    case class RefreshTokenSuccess(token: String)
    case class RefreshTokenError(error: String)

    case object Authenticate
    case object AuthenticationInProgress
    case class AuthenticationSuccess(token: String)
    case class AuthenticationError(error: String)

    case object FailedRecentlyRetryLater
  }
  case class State(token: Option[String], failedAt: Option[DateTime])
  private def now = DateTime.now(DateTimeZone.UTC)
}