Allow sending streams and not just receiving them (#121)

Author: ghostdogpr

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

@@ -32,10 +32,28 @@ trait Pods {
    */
   def sendMessage(pod: PodAddress, message: BinaryMessage): Task[Option[Array[Byte]]]
 
+  /**
+   * Send a stream of messages to a pod
+   */
+  def sendStream(
+    pod: PodAddress,
+    entityId: String,
+    messages: ZStream[Any, Throwable, BinaryMessage]
+  ): Task[Option[Array[Byte]]]
+
   /**
    * Send a message to a pod and receive a stream of replies
    */
-  def sendMessageStreaming(pod: PodAddress, message: BinaryMessage): ZStream[Any, Throwable, Array[Byte]]
+  def sendMessageAndReceiveStream(pod: PodAddress, message: BinaryMessage): ZStream[Any, Throwable, Array[Byte]]
+
+  /**
+   * Send a stream of messages to a pod and receive a stream of replies
+   */
+  def sendStreamAndReceiveStream(
+    pod: PodAddress,
+    entityId: String,
+    messages: ZStream[Any, Throwable, BinaryMessage]
+  ): ZStream[Any, Throwable, Array[Byte]]
 }
 
 object Pods {
@@ -46,12 +64,22 @@ object Pods {
    */
   val noop: ULayer[Pods] =
     ZLayer.succeed(new Pods {
-      def assignShards(pod: PodAddress, shards: Set[ShardId]): Task[Unit]                                     = ZIO.unit
-      def unassignShards(pod: PodAddress, shards: Set[ShardId]): Task[Unit]                                   = ZIO.unit
-      def ping(pod: PodAddress): Task[Unit]                                                                   = ZIO.unit
-      def sendMessage(pod: PodAddress, message: BinaryMessage): Task[Option[Array[Byte]]]                     = ZIO.none
-      def sendMessageStreaming(pod: PodAddress, message: BinaryMessage): ZStream[Any, Throwable, Array[Byte]] =
+      def assignShards(pod: PodAddress, shards: Set[ShardId]): Task[Unit]                                            = ZIO.unit
+      def unassignShards(pod: PodAddress, shards: Set[ShardId]): Task[Unit]                                          = ZIO.unit
+      def ping(pod: PodAddress): Task[Unit]                                                                          = ZIO.unit
+      def sendMessage(pod: PodAddress, message: BinaryMessage): Task[Option[Array[Byte]]]                            = ZIO.none
+      def sendStream(
+        pod: PodAddress,
+        entityId: String,
+        messages: ZStream[Any, Throwable, BinaryMessage]
+      ): 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
     })
 
   case class BinaryMessage(entityId: String, entityType: String, body: Array[Byte], replyId: Option[String])

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

@@ -27,7 +27,19 @@ trait Messenger[-Msg] {
    * streaming responses. See `sendStreamAutoRestart` for an alternative that will automatically restart the stream
    * in case of rebalance.
    */
-  def sendStream[Res](entityId: String)(msg: StreamReplier[Res] => Msg): Task[ZStream[Any, Throwable, Res]]
+  def sendAndReceiveStream[Res](entityId: String)(msg: StreamReplier[Res] => Msg): Task[ZStream[Any, Throwable, Res]]
+
+  /**
+   * Send a stream of messages.
+   */
+  def sendStream(entityId: String)(messages: ZStream[Any, Throwable, Msg]): Task[Unit]
+
+  /**
+   * Send a stream of messages and receive a stream of responses of type `Res`.
+   */
+  def sendStreamAndReceiveStream[Res](entityId: String)(
+    messages: StreamReplier[Res] => 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
@@ -38,11 +50,42 @@ trait Messenger[-Msg] {
    * cursor from the responses so that when the remote entity is rebalanced, a new message can be sent with the right
    * cursor according to what we've seen in the previous stream of responses.
    */
-  def sendStreamAutoRestart[Cursor, Res](entityId: String, cursor: Cursor)(msg: (Cursor, StreamReplier[Res]) => Msg)(
+  def sendAndReceiveStreamAutoRestart[Cursor, Res](entityId: String, cursor: Cursor)(
+    msg: (Cursor, StreamReplier[Res]) => Msg
+  )(
+    updateCursor: (Cursor, Res) => Cursor
+  ): ZStream[Any, Throwable, Res] =
+    ZStream
+      .unwrap(sendAndReceiveStream[Res](entityId)(msg(cursor, _)))
+      .either
+      .mapAccum(cursor) {
+        case (c, Right(res)) => updateCursor(c, res) -> Right(res)
+        case (c, Left(err))  => (c, Left(c -> err))
+      }
+      .flatMap {
+        case Right(res)                              => ZStream.succeed(res)
+        case Left((lastSeenCursor, StreamCancelled)) =>
+          ZStream.execute(ZIO.sleep(200.millis)) ++
+            sendAndReceiveStreamAutoRestart(entityId, lastSeenCursor)(msg)(updateCursor)
+        case Left((_, err))                          => ZStream.fail(err)
+      }
+
+  /**
+   * Send a stream of messages and receive a stream of responses of type `Res` while restarting the stream when the
+   * remote entity is rebalanced.
+   *
+   * To do so, we need a "cursor" so the stream of responses can be restarted where it ended before the rebalance. That
+   * is, the first message sent to the remote entity contains the given initial cursor value and we extract an updated
+   * cursor from the responses so that when the remote entity is rebalanced, a new message can be sent with the right
+   * 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]
+  )(
     updateCursor: (Cursor, Res) => Cursor
   ): ZStream[Any, Throwable, Res] =
     ZStream
-      .unwrap(sendStream[Res](entityId)(msg(cursor, _)))
+      .unwrap(sendStreamAndReceiveStream[Res](entityId)(msg(cursor, _)))
       .either
       .mapAccum(cursor) {
         case (c, Right(res)) => updateCursor(c, res) -> Right(res)
@@ -52,7 +95,7 @@ trait Messenger[-Msg] {
         case Right(res)                              => ZStream.succeed(res)
         case Left((lastSeenCursor, StreamCancelled)) =>
           ZStream.execute(ZIO.sleep(200.millis)) ++
-            sendStreamAutoRestart(entityId, lastSeenCursor)(msg)(updateCursor)
+            sendStreamAndReceiveStreamAutoRestart(entityId, lastSeenCursor)(msg)(updateCursor)
         case Left((_, err))                          => ZStream.fail(err)
       }
 }