[3.x] Add per-message MessageCodec backends with Protobuf support

.scalafmt.conf

@@ -1,4 +1,4 @@
-version = "3.1.2"
+version = "3.11.1"
 
 runner.dialect = scala3
 maxColumn = 120

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

@@ -1,5 +1,6 @@
 package com.devsisters.shardcake
 
+import com.devsisters.shardcake.kryo.given
 import com.devsisters.shardcake.interfaces.Storage
 import zio.{ Config => _, _ }
 import zio.stream.ZStream
@@ -58,7 +59,6 @@ object Server {
 
   val sharding: ZLayer[Config, Throwable, Sharding with GrpcConfig] =
     ZLayer.makeSome[Config, Sharding with GrpcConfig](
-      KryoSerialization.live,
       memory,
       grpcConfig,
       shardManagerClient,

build.sbt

@@ -14,7 +14,6 @@ val redis4catsVersion     = "2.0.1"
 val redissonVersion       = "3.45.1"
 val scalaKryoVersion      = "1.4.0"
 val testContainersVersion = "0.44.1"
-val scalaCompatVersion    = "2.13.0"
 
 inThisBuild(
   List(
@@ -54,6 +53,7 @@ lazy val root = project
     storageRedis,
     storageRedisson,
     serializationKryo,
+    serializationProteus,
     grpcProtocol,
     examples,
     benchmarks
@@ -66,10 +66,9 @@ lazy val core = project
   .settings(
     libraryDependencies ++=
       Seq(
-        "dev.zio"                %% "zio"                     % zioVersion,
-        "dev.zio"                %% "zio-streams"             % zioVersion,
-        "dev.zio"                %% "zio-json"                % zioJsonVersion,
-        "org.scala-lang.modules" %% "scala-collection-compat" % scalaCompatVersion
+        "dev.zio" %% "zio"         % zioVersion,
+        "dev.zio" %% "zio-streams" % zioVersion,
+        "dev.zio" %% "zio-json"    % zioJsonVersion
       )
   )
 
@@ -143,7 +142,7 @@ lazy val serializationKryo = project
   .in(file("serialization-kryo"))
   .settings(name := "shardcake-serialization-kryo")
   .settings(commonSettings)
-  .dependsOn(core)
+  .dependsOn(entities)
   .settings(
     libraryDependencies ++=
       Seq(
@@ -154,6 +153,18 @@ lazy val serializationKryo = project
 lazy val generateProto = taskKey[Unit]("Regenerate sharding.proto from the Scala protocol definition.")
 lazy val checkProto    = taskKey[Unit]("Fail if sharding.proto is out of sync with the Scala protocol definition.")
 
+lazy val serializationProteus = project
+  .in(file("serialization-proteus"))
+  .settings(name := "shardcake-serialization-proteus")
+  .settings(commonSettings)
+  .dependsOn(entities)
+  .settings(
+    libraryDependencies ++=
+      Seq(
+        "com.github.ghostdogpr" %% "proteus-core" % proteusVersion
+      )
+  )
+
 lazy val grpcProtocol = project
   .in(file("protocol-grpc"))
   .settings(name := "shardcake-protocol-grpc")

core/src/main/scala/com/devsisters/shardcake/PodAddress.scala

@@ -1,7 +1,5 @@
 package com.devsisters.shardcake
 
-import scala.collection.compat._
-
 import zio.json._
 
 case class PodAddress(host: String, port: Int) {

core/src/main/scala/com/devsisters/shardcake/interfaces/MessageCodec.scala

@@ -0,0 +1,51 @@
+package com.devsisters.shardcake.interfaces
+
+import com.devsisters.shardcake.interfaces.MessageCodec.{ Decoder, Encoder }
+
+/**
+ * Per-message-type codec used by Shardcake to serialise entity / topic messages and
+ * their replies on the wire.
+ *
+ * Carried by `RecipientType[Msg]` via an implicit `MessageCodec[Msg]` parameter, so each
+ * `EntityType` / `TopicType` declaration picks up the codec from the surrounding scope.
+ *
+ * Implementations decide how `Msg` is encoded and — for reply-bearing variants — which
+ * encoder/decoder pair to use for the reply type.
+ * For uniform backends (e.g. Kryo) the reply codec is the same regardless of which
+ * variant the message is; for variant-aware backends (e.g. Proteus) it depends on the
+ * concrete variant.
+ */
+trait MessageCodec[Msg] {
+
+  /**
+   * Encode a message to bytes for transport.
+   */
+  def encodeMessage(message: Msg): Array[Byte]
+
+  /**
+   * Decode a message from bytes.
+   */
+  def decodeMessage(bytes: Array[Byte]): Msg
+
+  /**
+   * Gets an encoder for the reply of a specific message on the receiver side.
+   * The returned function will be invoked when the entity calls `replier.reply(value)`
+   * or `streamReplier.replyStream(...)`.
+   */
+  def replyEncoder(decoded: Msg): Encoder
+
+  /**
+   * Returns a decoder for the reply of a specific message.
+   */
+  def replyDecoder[Res](sample: Msg): Decoder[Res]
+
+  /**
+   * Same as `replyDecoder`, but for stream-reply slots.
+   */
+  def streamReplyDecoder[Res](sample: Msg): Decoder[Res]
+}
+
+object MessageCodec {
+  type Encoder    = Any => Array[Byte]
+  type Decoder[A] = Array[Byte] => A
+}

core/src/main/scala/com/devsisters/shardcake/interfaces/Pods.scala

@@ -72,14 +72,14 @@ object Pods {
         pod: PodAddress,
         entityId: String,
         messages: ZStream[Any, Throwable, BinaryMessage]
-      ): Task[Option[Array[Byte]]] = ZIO.none
+      ): Task[Option[Array[Byte]]]                                                                                   = ZIO.none
       def sendMessageAndReceiveStream(pod: PodAddress, message: BinaryMessage): ZStream[Any, Throwable, Array[Byte]] =
         ZStream.empty
       def sendStreamAndReceiveStream(
         pod: PodAddress,
         entityId: String,
         messages: ZStream[Any, Throwable, BinaryMessage]
-      ): ZStream[Any, Throwable, Array[Byte]] = ZStream.empty
+      ): ZStream[Any, Throwable, Array[Byte]]                                                                        = ZStream.empty
     })
 
   case class BinaryMessage(entityId: String, entityType: String, body: Array[Byte], replyId: Option[String])

entities/src/main/scala/com/devsisters/shardcake/LocalSharding.scala

@@ -1,6 +1,6 @@
 package com.devsisters.shardcake
 
-import com.devsisters.shardcake.interfaces.{ Pods, Serialization, Storage }
+import com.devsisters.shardcake.interfaces.{ Pods, Storage }
 import com.devsisters.shardcake.interfaces.Pods.BinaryMessage
 import zio.{ Promise, Queue, RLayer, Task, ULayer, URLayer, ZIO, ZLayer }
 import zio.stream.ZStream
@@ -111,8 +111,8 @@ object LocalSharding {
    * A special layer meant for testing that uses a local queue rather than an external transport.
    * This layer will only work in a single JVM and is not suitable for production use.
    */
-  val live: RLayer[ShardManagerClient with Storage with Serialization with Config, Sharding] =
-    ZLayer.makeSome[ShardManagerClient with Storage with Serialization with Config, Sharding](
+  val live: RLayer[ShardManagerClient with Storage with Config, Sharding] =
+    ZLayer.makeSome[ShardManagerClient with Storage with Config, Sharding](
       localQueue,
       localPods,
       localServer,

entities/src/main/scala/com/devsisters/shardcake/Messenger.scala

@@ -36,10 +36,13 @@ trait Messenger[-Msg] {
 
   /**
    * Send a stream of messages and receive a stream of responses of type `Res`.
+   *
+   * The first message carrying the `StreamReplier` is constructed via `request`; any
+   * additional messages flow through `rest`.
    */
-  def sendStreamAndReceiveStream[Res](entityId: String)(
-    messages: StreamReplier[Res] => ZStream[Any, Throwable, Msg]
-  ): Task[ZStream[Any, Throwable, Res]]
+  def sendStreamAndReceiveStream[Res](
+    entityId: String
+  )(request: StreamReplier[Res] => Msg, rest: ZStream[Any, Throwable, Msg]): Task[ZStream[Any, Throwable, Res]]
 
   /**
    * Send a message and receive a stream of responses of type `Res` while restarting the stream when the remote entity
@@ -80,12 +83,13 @@ trait Messenger[-Msg] {
    * cursor according to what we've seen in the previous stream of responses.
    */
   def sendStreamAndReceiveStreamAutoRestart[Cursor, Res](entityId: String, cursor: Cursor)(
-    msg: (Cursor, StreamReplier[Res]) => ZStream[Any, Throwable, Msg]
+    request: (Cursor, StreamReplier[Res]) => Msg,
+    rest: Cursor => ZStream[Any, Throwable, Msg]
   )(
     updateCursor: (Cursor, Res) => Cursor
   ): ZStream[Any, Throwable, Res] =
     ZStream
-      .unwrap(sendStreamAndReceiveStream[Res](entityId)(msg(cursor, _)))
+      .unwrap(sendStreamAndReceiveStream[Res](entityId)(request(cursor, _), rest(cursor)))
       .either
       .mapAccum(cursor) {
         case (c, Right(res)) => updateCursor(c, res) -> Right(res)
@@ -95,7 +99,7 @@ trait Messenger[-Msg] {
         case Right(res)                              => ZStream.succeed(res)
         case Left((lastSeenCursor, StreamCancelled)) =>
           ZStream.execute(ZIO.sleep(200.millis)) ++
-            sendStreamAndReceiveStreamAutoRestart(entityId, lastSeenCursor)(msg)(updateCursor)
+            sendStreamAndReceiveStreamAutoRestart(entityId, lastSeenCursor)(request, rest)(updateCursor)
         case Left((_, err))                          => ZStream.fail(err)
       }
 }

entities/src/main/scala/com/devsisters/shardcake/RecipientType.scala

@@ -1,13 +1,20 @@
 package com.devsisters.shardcake
 
+import com.devsisters.shardcake.interfaces.MessageCodec
+
 /**
- * An abstract type to extend for each type of entity or topic
+ * An abstract type to extend for each type of entity or topic.
+ *
+ * The `MessageCodec[Msg]` carried via the implicit parameter owns serialisation for this
+ * recipient type. Users provide a codec by importing one of the backend `given`s in scope
+ * at the declaration site, e.g. `import com.devsisters.shardcake.kryo.given`.
+ *
  * @param name a unique string that identifies this entity or topic type
  * @tparam Msg the type of message that can be sent to this entity or topic type
  */
-sealed abstract class RecipientType[+Msg](val name: String) {
+sealed abstract class RecipientType[Msg](val name: String)(using private[shardcake] val codec: MessageCodec[Msg]) {
   def getShardId(entityId: String, numberOfShards: Int): ShardId =
     math.abs(entityId.hashCode % numberOfShards) + 1
 }
-abstract class EntityType[+Msg](name: String) extends RecipientType[Msg](name)
-abstract class TopicType[+Msg](name: String) extends RecipientType[Msg](name)
+abstract class EntityType[Msg](name: String)(using MessageCodec[Msg]) extends RecipientType[Msg](name)
+abstract class TopicType[Msg](name: String)(using MessageCodec[Msg]) extends RecipientType[Msg](name)