akka.cluster.sharding.ShardRegion Scala Examples

package aia.persistence.sharded

import aia.persistence._
import akka.actor._
import akka.cluster.sharding.ShardRegion
import akka.cluster.sharding.ShardRegion.Passivate

object ShardedShopper {
  def props = Props(new ShardedShopper)
  def name(shopperId: Long) = shopperId.toString


  case object StopShopping

  val shardName: String = "shoppers"

  val extractEntityId: ShardRegion.ExtractEntityId = {
    case cmd: Shopper.Command => (cmd.shopperId.toString, cmd)
  }

  val extractShardId: ShardRegion.ExtractShardId = {
    case cmd: Shopper.Command => (cmd.shopperId % 12).toString
  }
}

class ShardedShopper extends Shopper {
  import ShardedShopper._

  context.setReceiveTimeout(Settings(context.system).passivateTimeout)

  override def unhandled(msg: Any) = msg match {
    case ReceiveTimeout =>
      context.parent ! Passivate(stopMessage = ShardedShopper.StopShopping)
    case StopShopping => context.stop(self)
  }
} 

package akka.persistence.cassandra.sharding

import akka.actor.{ ActorLogging, ActorRef, Props, ReceiveTimeout }
import akka.cluster.{ Cluster, MemberStatus }
import akka.cluster.sharding.{ ClusterSharding, ClusterShardingSettings, ShardRegion }
import akka.persistence.PersistentActor
import akka.persistence.cassandra.CassandraSpec
import akka.testkit.TestProbe

import scala.concurrent.duration._

object ClusterShardingQuickTerminationSpec {

  case object Increment
  case object Decrement
  final case class Get(counterId: Long)
  final case class EntityEnvelope(id: Long, payload: Any)
  case object Ack

  case object Stop
  final case class CounterChanged(delta: Int)

  class Counter extends PersistentActor with ActorLogging {
    import ShardRegion.Passivate

    context.setReceiveTimeout(5.seconds)

    // self.path.name is the entity identifier (utf-8 URL-encoded)
    override def persistenceId: String = "Counter-" + self.path.name

    var count = 0

    def updateState(event: CounterChanged): Unit =
      count += event.delta

    override def receiveRecover: Receive = {
      case evt: CounterChanged => updateState(evt)
      case other               => log.debug("Other: {}", other)
    }

    override def receiveCommand: Receive = {
      case Increment      => persist(CounterChanged(+1))(updateState)
      case Decrement      => persist(CounterChanged(-1))(updateState)
      case Get(_)         => sender() ! count
      case ReceiveTimeout => context.parent ! Passivate(stopMessage = Stop)
      case Stop =>
        sender() ! Ack
        context.stop(self)
    }
  }
  val extractEntityId: ShardRegion.ExtractEntityId = {
    case EntityEnvelope(id, payload) => (id.toString, payload)
    case msg @ Get(id)               => (id.toString, msg)
  }

  val numberOfShards = 100

  val extractShardId: ShardRegion.ExtractShardId = {
    case EntityEnvelope(id, _) => (id % numberOfShards).toString
    case Get(id)               => (id % numberOfShards).toString
  }

}

class ClusterShardingQuickTerminationSpec extends CassandraSpec("""
    akka.actor.provider = cluster
  """.stripMargin) {

  import ClusterShardingQuickTerminationSpec._

  "Cassandra Plugin with Cluster Sharding" must {
    "clear state if persistent actor shuts down" in {
      Cluster(system).join(Cluster(system).selfMember.address)
      awaitAssert {
        Cluster(system).selfMember.status shouldEqual MemberStatus.Up
      }
      ClusterSharding(system).start(
        typeName = "tagging",
        entityProps = Props[Counter],
        settings = ClusterShardingSettings(system),
        extractEntityId = extractEntityId,
        extractShardId = extractShardId)

      (0 to 100).foreach { i =>
        val counterRegion: ActorRef = ClusterSharding(system).shardRegion("tagging")
        awaitAssert {
          val sender = TestProbe()
          counterRegion.tell(Get(123), sender.ref)
          sender.expectMsg(500.millis, i)
        }

        counterRegion ! EntityEnvelope(123, Increment)
        counterRegion ! Get(123)
        expectMsg(i + 1)

        counterRegion ! EntityEnvelope(123, Stop)
        expectMsg(Ack)
      }
    }
  }
} 

package aecor.distributedprocessing

import aecor.distributedprocessing.DistributedProcessingSupervisor.{
  GracefulShutdown,
  ShutdownCompleted,
  Tick
}
import aecor.distributedprocessing.DistributedProcessingWorker.KeepRunning
import akka.actor.{ Actor, ActorLogging, ActorRef, Props, Terminated }
import akka.cluster.sharding.ShardRegion

import scala.concurrent.duration.{ FiniteDuration, _ }

object DistributedProcessingSupervisor {
  private final case object Tick
  final case object GracefulShutdown
  final case object ShutdownCompleted

  def props(processCount: Int, shardRegion: ActorRef, heartbeatInterval: FiniteDuration): Props =
    Props(new DistributedProcessingSupervisor(processCount, shardRegion, heartbeatInterval))
}

final class DistributedProcessingSupervisor(processCount: Int,
                                            shardRegion: ActorRef,
                                            heartbeatInterval: FiniteDuration)
    extends Actor
    with ActorLogging {

  import context.dispatcher

  private val heartbeat =
    context.system.scheduler.schedule(0.seconds, heartbeatInterval, self, Tick)

  context.watch(shardRegion)

  override def postStop(): Unit = {
    heartbeat.cancel()
    ()
  }

  override def receive: Receive = {
    case Tick =>
      (0 until processCount).foreach { processId =>
        shardRegion ! KeepRunning(processId)
      }
    case Terminated(`shardRegion`) =>
      context.stop(self)
    case GracefulShutdown =>
      log.info(s"Performing graceful shutdown of [$shardRegion]")
      shardRegion ! ShardRegion.GracefulShutdown
      val replyTo = sender()
      context.become {
        case Terminated(`shardRegion`) =>
          log.info(s"Graceful shutdown completed for [$shardRegion]")
          context.stop(self)
          replyTo ! ShutdownCompleted
      }

  }
} 

package aecor.runtime.akkageneric

import aecor.encoding.WireProtocol.Encoded
import aecor.encoding.syntax._
import aecor.encoding.{ KeyDecoder, KeyEncoder, WireProtocol }
import aecor.runtime.akkageneric.GenericAkkaRuntime.KeyedCommand
import aecor.runtime.akkageneric.GenericAkkaRuntimeActor.CommandResult
import aecor.runtime.akkageneric.serialization.Message
import aecor.util.effect._
import akka.actor.ActorSystem
import akka.cluster.sharding.{ ClusterSharding, ShardRegion }
import akka.pattern._
import akka.util.Timeout
import cats.effect.Effect
import cats.implicits._
import cats.tagless.FunctorK
import cats.tagless.syntax.functorK._
import cats.~>
import scodec.bits.BitVector

object GenericAkkaRuntime {
  def apply(system: ActorSystem): GenericAkkaRuntime =
    new GenericAkkaRuntime(system)
  private[akkageneric] final case class KeyedCommand(key: String, bytes: BitVector) extends Message
}

final class GenericAkkaRuntime private (system: ActorSystem) {
  def runBehavior[K: KeyEncoder: KeyDecoder, M[_[_]]: FunctorK, F[_]](
    typeName: String,
    createBehavior: K => F[M[F]],
    settings: GenericAkkaRuntimeSettings = GenericAkkaRuntimeSettings.default(system)
  )(implicit M: WireProtocol[M], F: Effect[F]): F[K => M[F]] =
    F.delay {
      val props = GenericAkkaRuntimeActor.props[K, M, F](createBehavior, settings.idleTimeout)

      val extractEntityId: ShardRegion.ExtractEntityId = {
        case KeyedCommand(entityId, c) =>
          (entityId, GenericAkkaRuntimeActor.Command(c))
      }

      val numberOfShards = settings.numberOfShards

      val extractShardId: ShardRegion.ExtractShardId = {
        case KeyedCommand(key, _) =>
          String.valueOf(scala.math.abs(key.hashCode) % numberOfShards)
        case other => throw new IllegalArgumentException(s"Unexpected message [$other]")
      }

      val shardRegion = ClusterSharding(system).start(
        typeName = typeName,
        entityProps = props,
        settings = settings.clusterShardingSettings,
        extractEntityId = extractEntityId,
        extractShardId = extractShardId
      )

      val keyEncoder = KeyEncoder[K]

      key =>
        M.encoder.mapK(new (Encoded ~> F) {

          implicit val askTimeout: Timeout = Timeout(settings.askTimeout)

          override def apply[A](fa: Encoded[A]): F[A] = F.suspend {
            val (bytes, decoder) = fa
            F.fromFuture {
                shardRegion ? KeyedCommand(keyEncoder(key), bytes)
              }
              .flatMap {
                case result: CommandResult =>
                  decoder.decodeValue(result.bytes).lift[F]
                case other =>
                  F.raiseError(
                    new IllegalArgumentException(s"Unexpected response [$other] from shard region")
                  )
              }
          }
        })
    }
} 

package aecor.runtime.akkapersistence

import aecor.data.{ EventsourcedBehavior, Tagging }
import aecor.encoding.WireProtocol.Encoded
import aecor.encoding.syntax._
import aecor.encoding.{ KeyDecoder, KeyEncoder, WireProtocol }
import aecor.runtime.akkapersistence.AkkaPersistenceRuntime._
import aecor.runtime.akkapersistence.AkkaPersistenceRuntimeActor.CommandResult
import aecor.runtime.akkapersistence.readside.{ AkkaPersistenceEventJournalQuery, JournalQuery }
import aecor.runtime.akkapersistence.serialization.{ Message, PersistentDecoder, PersistentEncoder }
import aecor.util.effect._
import akka.actor.ActorSystem
import akka.cluster.sharding.{ ClusterSharding, ShardRegion }
import akka.pattern.ask
import akka.util.Timeout
import cats.effect.Effect
import cats.implicits._
import cats.tagless.FunctorK
import cats.tagless.syntax.functorK._
import cats.~>
import scodec.bits.BitVector

object AkkaPersistenceRuntime {
  def apply[O](system: ActorSystem, journalAdapter: JournalAdapter[O]): AkkaPersistenceRuntime[O] =
    new AkkaPersistenceRuntime(system, journalAdapter)

  private[akkapersistence] final case class EntityCommand(entityKey: String,
                                                          commandBytes: BitVector)
      extends Message
}

class AkkaPersistenceRuntime[O] private[akkapersistence] (system: ActorSystem,
                                                          journalAdapter: JournalAdapter[O]) {
  def deploy[M[_[_]]: FunctorK, F[_], State, Event: PersistentEncoder: PersistentDecoder, K: KeyEncoder: KeyDecoder](
    typeName: String,
    behavior: EventsourcedBehavior[M, F, State, Event],
    tagging: Tagging[K],
    snapshotPolicy: SnapshotPolicy[State] = SnapshotPolicy.never,
    settings: AkkaPersistenceRuntimeSettings = AkkaPersistenceRuntimeSettings.default(system)
  )(implicit M: WireProtocol[M], F: Effect[F]): F[K => M[F]] =
    F.delay {
      val props =
        AkkaPersistenceRuntimeActor.props(
          typeName,
          behavior,
          snapshotPolicy,
          tagging,
          settings.idleTimeout,
          journalAdapter.writeJournalId,
          snapshotPolicy.pluginId
        )

      val extractEntityId: ShardRegion.ExtractEntityId = {
        case EntityCommand(entityId, bytes) =>
          (entityId, AkkaPersistenceRuntimeActor.HandleCommand(bytes))
      }

      val numberOfShards = settings.numberOfShards

      val extractShardId: ShardRegion.ExtractShardId = {
        case EntityCommand(entityId, _) =>
          (scala.math.abs(entityId.hashCode) % numberOfShards).toString
        case other => throw new IllegalArgumentException(s"Unexpected message [$other]")
      }

      val shardRegion = ClusterSharding(system).start(
        typeName = typeName,
        entityProps = props,
        settings = settings.clusterShardingSettings,
        extractEntityId = extractEntityId,
        extractShardId = extractShardId
      )

      val keyEncoder = KeyEncoder[K]

      key =>
        M.encoder.mapK(new (Encoded ~> F) {

          implicit val askTimeout: Timeout = Timeout(settings.askTimeout)

          override def apply[A](fa: Encoded[A]): F[A] = F.suspend {
            val (bytes, decoder) = fa
            F.fromFuture {
                shardRegion ? EntityCommand(keyEncoder(key), bytes)
              }
              .flatMap {
                case CommandResult(resultBytes) =>
                  decoder.decodeValue(resultBytes).lift[F]
                case other =>
                  F.raiseError(
                    new IllegalArgumentException(s"Unexpected response [$other] from shard region")
                  )
              }
          }
        })
    }

  def journal[K: KeyDecoder, E: PersistentDecoder]: JournalQuery[O, K, E] =
    AkkaPersistenceEventJournalQuery[O, K, E](journalAdapter)
} 

package cluster

import java.util.UUID

import akka.actor.SupervisorStrategy.Stop
import akka.actor.{ Actor, ActorLogging, Props, ReceiveTimeout }
import akka.cluster.sharding.ShardRegion
import akka.cluster.sharding.ShardRegion.Passivate
import cluster.CacheDataActor.Get

class CacheDataActor extends Actor with ActorLogging {

  override def receive: Receive = {
    case Get(id) => sender ! s"cached data for id: $id"
    case ReceiveTimeout =>
      log.info(s"sending Passivate to metadata parent: {${context.parent.path.name}} for ${self.path.name}")
      context.parent ! Passivate(stopMessage = Stop)
    case Stop =>
      context.stop(self)
      log.info(s"Passivating metadata actor for ${self.path.name}")
  }
}

object CacheDataActor {
  final val numOfShards = 50 // Planned num of cluster nodes
  val extractEntityId: ShardRegion.ExtractEntityId = {
    case msg@Get(id) => (id.toString, msg)
  }
  val extractShardId: ShardRegion.ExtractShardId = {
    case Get(id) => (id.hashCode() % numOfShards).toString
  }

  case class Get(id: UUID)

  def props: Props = Props(new CacheDataActor())
} 

package connectedcar.actors

import akka.actor.{ Actor, ActorLogging, ActorRef }
import akka.cluster.sharding.ShardRegion
import connectedcar.data.{ ConnectedCarAgg, ConnectedCarERecord }

import scala.concurrent.ExecutionContext

object ConnectedCarActor {
  val extractEntityId: ShardRegion.ExtractEntityId = {
    case msg: ConnectedCarERecord ⇒ (msg.carId.toString, msg)
  }

  private val numberOfShards = 100

  val extractShardId: ShardRegion.ExtractShardId = {
    case msg: ConnectedCarERecord ⇒ (msg.carId % numberOfShards).toString
  }
}

class ConnectedCarActor extends Actor with ActorLogging {

  val carId: String      = "Car-" + self.path.name
  var driverName: String = null
  var currentSpeed       = 0.0
  var averageSpeed       = 0.0
  var numberOfRecords    = 0

  var treeActor: ActorRef           = null
  implicit val ec: ExecutionContext = context.dispatcher

  override def receive: Receive = {
    case record: ConnectedCarERecord ⇒ {
      if (numberOfRecords == 0) {
        driverName = record.driver
        averageSpeed = record.speed
      } else {
        averageSpeed = ((averageSpeed * numberOfRecords) + record.speed) / (numberOfRecords + 1)
      }

      numberOfRecords += 1
      currentSpeed = record.speed

      log.info("Updated CarId: " + carId + " Driver Name: " + driverName + " CarSpeed: " + currentSpeed + " From Actor:" + sender().path)

      sender() ! ConnectedCarAgg(record.carId, record.driver, averageSpeed, numberOfRecords)
    }
  }
} 

package akka.kubernetes.sample

import akka.actor.{Actor, Stash}
import akka.cluster.sharding.ShardRegion
import akka.cluster.singleton.{ClusterSingletonProxy, ClusterSingletonProxySettings}
import akka.event.Logging
import akka.kubernetes.sample.AkkaBoss.{GoToJobCentre, JobSpec, WhatCanIDo}
import akka.kubernetes.sample.AkkaMember.Hello

case object AkkaMember {
  case class Hello(name: String)

  val extractEntityId: ShardRegion.ExtractEntityId = {
    case msg @ Hello(id) ⇒ (id, msg)
  }

  val numberOfShards = 100

  val extractShardId: ShardRegion.ExtractShardId = {
    case Hello(id) ⇒ math.abs(id.hashCode % numberOfShards).toString
  }
}

class AkkaMember() extends Actor with Stash {
  val bossProxy = context.system.actorOf(
    ClusterSingletonProxy.props(singletonManagerPath = "/user/boss",
                                settings = ClusterSingletonProxySettings(context.system))
  )

  val name = self.path.name
  val log = Logging(this)

  override def preStart(): Unit = {
    // TODO retry
    log.info("/me good morning {}", name)
    bossProxy ! WhatCanIDo(name)
  }

  override def receive: Receive = {
    case JobSpec(roles) =>
      log.info("I'm part of the team. I can do {}", roles)
      unstashAll()
      context.become(ready(roles))
    case GoToJobCentre(n) =>
      log.info("Seems I am not in the team :( I'll go fruit picking")
      context.stop(self)
    case _ =>
      stash()
  }

  def ready(roles: Set[String]): Receive = {
    case Hello(_) =>
      sender() ! s"hello from $name. What can I do for you? It better be in ${roles.mkString(", ")}"
    case _ =>
      sender() ! "what?"
  }
} 

package com.michalplachta.shoesorter

import akka.actor.{ActorLogging, Actor, Props}
import akka.cluster.sharding.ShardRegion
import akka.cluster.sharding.ShardRegion.{ExtractEntityId, ExtractShardId}
import com.michalplachta.shoesorter.Domain.{Container, Junction}
import com.michalplachta.shoesorter.Messages._

object SortingDecider {
  def name = "sortingDecider"

  def props = Props[SortingDecider]

  def extractShardId: ExtractShardId = {
    case WhereShouldIGo(junction, _) =>
      (junction.id % 2).toString
  }

  def extractEntityId: ExtractEntityId = {
    case msg @ WhereShouldIGo(junction, _) =>
      (junction.id.toString, msg)
  }
}

class SortingDecider extends Actor with ActorLogging {
  def receive = {
    case WhereShouldIGo(junction, container) =>
      val decision = Decisions.whereShouldContainerGo(junction, container)
      log.info("Decision on junction {} for container {}: {}", junction.id, container.id, decision)
      sender ! Go(decision)
  }
}