Add benchmarks module (#128)

Author: ghostdogpr

benchmarks/src/main/scala/com/devsisters/shardcake/Client.scala

@@ -0,0 +1,21 @@
+package com.devsisters.shardcake
+
+import com.devsisters.shardcake.Config
+import com.devsisters.shardcake.Server.Message.Ping
+import com.devsisters.shardcake.Server.PingPongEntity
+import zio._
+
+object Client {
+  // self host should not be `localhost` to avoid optimization
+  private val config = ZLayer.succeed(Config.default.copy(selfHost = "not-localhost"))
+
+  def send(count: Int, parallelism: Int): Task[Unit] =
+    ZIO
+      .scoped[Sharding](
+        for {
+          ping <- Sharding.messenger(PingPongEntity)
+          _    <- ZIO.foreachParDiscard(1 to count)(_ => ping.send("ping")(Ping("ping", _))).withParallelism(parallelism)
+        } yield ()
+      )
+      .provide(config, Server.sharding)
+}

benchmarks/src/main/scala/com/devsisters/shardcake/SendBenchmark.scala

@@ -0,0 +1,30 @@
+package com.devsisters.shardcake
+
+import org.openjdk.jmh.annotations._
+import zio.{ durationInt, Fiber, Runtime, Unsafe, ZIO }
+
+import java.util.concurrent.TimeUnit
+
+@State(Scope.Thread)
+@BenchmarkMode(Array(Mode.Throughput))
+@OutputTimeUnit(TimeUnit.SECONDS)
+@Warmup(iterations = 5, time = 5, timeUnit = TimeUnit.SECONDS)
+@Measurement(iterations = 5, time = 5, timeUnit = TimeUnit.SECONDS)
+@Fork(1)
+class SendBenchmark {
+  private var fiber: Fiber[Any, Any] = _
+
+  @Setup
+  def setup(): Unit =
+    fiber = Unsafe.unsafe(implicit unsafe =>
+      Runtime.default.unsafe.run(Server.run.forkDaemon <* ZIO.sleep(3.seconds)).getOrThrow()
+    )
+
+  @TearDown
+  def tearDown(): Unit =
+    Unsafe.unsafe(implicit unsafe => Runtime.default.unsafe.run(fiber.interrupt))
+
+  @Benchmark
+  def send(): Unit =
+    Unsafe.unsafe(implicit unsafe => Runtime.default.unsafe.run(Client.send(100, 8)))
+}

benchmarks/src/main/scala/com/devsisters/shardcake/Server.scala

@@ -0,0 +1,77 @@
+package com.devsisters.shardcake
+
+import com.devsisters.shardcake.Config
+import com.devsisters.shardcake.interfaces.Storage
+import zio._
+import zio.stream.ZStream
+
+import java.util.concurrent.ForkJoinPool
+
+object Server {
+  sealed trait Message
+
+  object Message {
+    case class Ping(msg: String, replier: Replier[String]) extends Message
+  }
+
+  object PingPongEntity extends EntityType[Message]("ping-pong")
+
+  private def behavior(entityId: String, messages: Dequeue[Message]): RIO[Sharding, Nothing] =
+    messages.take.flatMap { case Message.Ping(msg, replier) => replier.reply(msg) }.forever
+
+  private val shardManagerClient: ZLayer[Config, Nothing, ShardManagerClient] =
+    ZLayer {
+      for {
+        config <- ZIO.service[Config]
+        pod     = PodAddress("localhost", config.shardingPort)
+        shards  = (1 to config.numberOfShards).map(_ -> Some(pod)).toMap
+      } yield new ShardManagerClient {
+        def register(podAddress: PodAddress): Task[Unit]           = ZIO.unit
+        def unregister(podAddress: PodAddress): Task[Unit]         = ZIO.unit
+        def notifyUnhealthyPod(podAddress: PodAddress): Task[Unit] = ZIO.unit
+        def getAssignments: Task[Map[Int, Option[PodAddress]]]     = ZIO.succeed(shards)
+      }
+    }
+
+  private val memory: ULayer[Storage] =
+    ZLayer {
+      for {
+        assignmentsRef <- Ref.make(Map.empty[ShardId, Option[PodAddress]])
+        podsRef        <- Ref.make(Map.empty[PodAddress, Pod])
+      } yield new Storage {
+        def getAssignments: Task[Map[ShardId, Option[PodAddress]]]                       = assignmentsRef.get
+        def saveAssignments(assignments: Map[ShardId, Option[PodAddress]]): Task[Unit]   = assignmentsRef.set(assignments)
+        def assignmentsStream: ZStream[Any, Throwable, Map[ShardId, Option[PodAddress]]] = ZStream.never
+        def getPods: Task[Map[PodAddress, Pod]]                                          = podsRef.get
+        def savePods(pods: Map[PodAddress, Pod]): Task[Unit]                             = podsRef.set(pods)
+      }
+    }
+
+  private val grpcConfig: ULayer[GrpcConfig] =
+    ZLayer.succeed(GrpcConfig.default.copy(executor = Some(ForkJoinPool.commonPool())))
+
+  val sharding: ZLayer[Config, Throwable, Sharding with GrpcConfig] =
+    ZLayer.makeSome[Config, Sharding with GrpcConfig](
+      KryoSerialization.live,
+      memory,
+      grpcConfig,
+      shardManagerClient,
+      GrpcPods.live,
+      Sharding.live
+    )
+
+  val run: Task[Unit] =
+    ZIO
+      .scoped[Sharding](
+        for {
+          _ <- Sharding.registerEntity(PingPongEntity, behavior)
+          _ <- Sharding.registerScoped
+          _ <- ZIO.never
+        } yield ()
+      )
+      .provide(
+        ZLayer.succeed(Config.default),
+        sharding,
+        GrpcShardingService.live
+      )
+}

build.sbt

@@ -58,7 +58,8 @@ lazy val root = project
     storageRedisson,
     serializationKryo,
     grpcProtocol,
-    examples
+    examples,
+    benchmarks
   )
 
 lazy val core = project
@@ -186,6 +187,14 @@ lazy val examples = project
   )
   .dependsOn(manager, storageRedis, grpcProtocol, serializationKryo)
 
+lazy val benchmarks = project
+  .in(file("benchmarks"))
+  .settings(name := "benchmarks")
+  .settings(publish / skip := true)
+  .settings(commonSettings)
+  .enablePlugins(JmhPlugin)
+  .dependsOn(grpcProtocol, serializationKryo)
+
 lazy val protobuf = Seq(
   PB.protocVersion := "3.19.2"
 ) ++ Project.inConfig(Test)(sbtprotoc.ProtocPlugin.protobufConfigSettings)

project/plugins.sbt

@@ -1,7 +1,8 @@
-addSbtPlugin("org.scalameta"  % "sbt-scalafmt"   % "2.4.6")
-addSbtPlugin("com.github.sbt" % "sbt-ci-release" % "1.5.10")
-addSbtPlugin("org.xerial.sbt" % "sbt-sonatype"   % "3.9.13")
-addSbtPlugin("com.thesamet"   % "sbt-protoc"     % "1.0.6")
+addSbtPlugin("org.scalameta"      % "sbt-scalafmt"   % "2.4.6")
+addSbtPlugin("com.github.sbt"     % "sbt-ci-release" % "1.5.10")
+addSbtPlugin("org.xerial.sbt"     % "sbt-sonatype"   % "3.9.13")
+addSbtPlugin("com.thesamet"       % "sbt-protoc"     % "1.0.6")
+addSbtPlugin("pl.project13.scala" % "sbt-jmh"        % "0.4.7")
 
 resolvers ++= Resolver.sonatypeOssRepos("snapshots")