Fix client logic for backpressuring streams with prefetch

Author: joroKr21

core/src/main/scalajvm/scalapb/zio_grpc/client/StreamingClientCallListener.scala

@@ -9,34 +9,36 @@ class StreamingClientCallListener[Res](
     prefetch: Option[Int],
     runtime: Runtime[Any],
     call: ZClientCall[?, Res],
-    queue: Queue[ResponseFrame[Res]],
-    buffered: Ref[Int]
+    queue: Queue[ResponseFrame[Res]]
 ) extends ClientCall.Listener[Res] {
-  private val increment = if (prefetch.isDefined) buffered.update(_ + 1) else ZIO.unit
-  private val fetchOne  = if (prefetch.isDefined) ZIO.unit else call.request(1)
-  private val fetchMore = prefetch match {
-    case None    => ZIO.unit
-    case Some(n) => buffered.get.flatMap(b => call.request(n - b).when(n > b))
-  }
+  private val fetchOne =
+    if (prefetch.isDefined) ZIO.unit else call.request(1)
+
+  private def fetchMore(n: Int) =
+    if (prefetch.isDefined) call.request(n) else ZIO.unit
 
   private def unsafeRun(task: IO[Any, Unit]): Unit =
     Unsafe.unsafe(implicit u => runtime.unsafe.run(task).getOrThrowFiberFailure())
 
-  private def handle(promise: Promise[StatusException, Unit])(
-      chunk: Chunk[ResponseFrame[Res]]
-  ) = (chunk.lastOption match {
-    case Some(ResponseFrame.Trailers(status, trailers)) =>
-      val exit = if (status.isOk) Exit.unit else Exit.fail(new StatusException(status, trailers))
-      promise.done(exit) *> queue.shutdown
-    case _                                              =>
-      buffered.update(_ - chunk.size) *> fetchMore
-  }).as(chunk)
+  private def handle(promise: Promise[StatusException, Unit])(chunk: Chunk[ResponseFrame[Res]]) =
+    if (chunk.isEmpty) ZIO.unit
+    else {
+      chunk.last match {
+        case ResponseFrame.Trailers(status, trailers) =>
+          val exit =
+            if (status.isOk) Exit.unit
+            else Exit.fail(new StatusException(status, trailers))
+          promise.done(exit) *> queue.shutdown
+        case _                                        =>
+          fetchMore(chunk.size)
+      }
+    }
 
   override def onHeaders(headers: Metadata): Unit =
-    unsafeRun(queue.offer(ResponseFrame.Headers(headers)) *> increment)
+    unsafeRun(queue.offer(ResponseFrame.Headers(headers)).unit)
 
   override def onMessage(message: Res): Unit =
-    unsafeRun(queue.offer(ResponseFrame.Message(message)) *> increment *> fetchOne)
+    unsafeRun(queue.offer(ResponseFrame.Message(message)) *> fetchOne)
 
   override def onClose(status: Status, trailers: Metadata): Unit =
     unsafeRun(queue.offer(ResponseFrame.Trailers(status, trailers)).unit)
@@ -45,15 +47,14 @@ class StreamingClientCallListener[Res](
     ZStream.fromZIO(Promise.make[StatusException, Unit]).flatMap { promise =>
       ZStream
         .fromQueue(queue, prefetch.getOrElse(ZStream.DefaultChunkSize))
-        .mapChunksZIO(handle(promise))
+        .tapChunks(handle(promise))
         .concat(ZStream.execute(promise.await))
     }
 }
 
 object StreamingClientCallListener {
   def make[Res](call: ZClientCall[?, Res], prefetch: Option[Int]): UIO[StreamingClientCallListener[Res]] = for {
-    runtime  <- ZIO.runtime[Any]
-    queue    <- Queue.unbounded[ResponseFrame[Res]]
-    buffered <- Ref.make(0)
-  } yield new StreamingClientCallListener(prefetch, runtime, call, queue, buffered)
+    runtime <- ZIO.runtime[Any]
+    queue   <- Queue.unbounded[ResponseFrame[Res]]
+  } yield new StreamingClientCallListener(prefetch, runtime, call, queue)
 }

e2e/protos/src/main/protobuf/testservice.proto

@@ -5,38 +5,37 @@ package scalapb.zio_grpc;
 import "scalapb/scalapb.proto";
 
 message Request {
-    enum Scenario {
-        OK = 0;
-        ERROR_NOW = 1; // fail with an error
-        ERROR_AFTER = 2; // for server streaming, error after two responses
-        DELAY = 3; // do not return a response. for testing cancellations
-        DIE = 4; // fail
-        UNAVAILABLE = 5; // fail with UNAVAILABLE, to test client retries
-    }
-    Scenario scenario = 1;
-    int32 in = 2;
+  enum Scenario {
+    OK = 0;
+    ERROR_NOW = 1;    // fail with an error
+    ERROR_AFTER = 2;  // for server streaming, error after two responses
+    DELAY = 3;        // do not return a response, to test cancellations
+    LARGE_STREAM = 4; // stream of large elements, to test backpressure
+    DIE = 5;          // fail
+    UNAVAILABLE = 6;  // fail with UNAVAILABLE, to test client retries
+  }
+  Scenario scenario = 1;
+  int32 in = 2;
 }
 
-message Response {
-    string out = 1;
-}
+message Response { string out = 1; }
 
 message ResponseTypeMapped {
-    option (scalapb.message).type = "scalapb.zio_grpc.WrappedString";
+  option (scalapb.message).type = "scalapb.zio_grpc.WrappedString";
 
-    string out = 1;
+  string out = 1;
 }
 
 service TestService {
-    rpc Unary(Request) returns (Response);
+  rpc Unary(Request) returns (Response);
 
-    rpc UnaryTypeMapped(Request) returns (ResponseTypeMapped);
+  rpc UnaryTypeMapped(Request) returns (ResponseTypeMapped);
 
-    rpc ServerStreaming(Request) returns (stream Response);
+  rpc ServerStreaming(Request) returns (stream Response);
 
-    rpc ServerStreamingTypeMapped(Request) returns (stream ResponseTypeMapped);
+  rpc ServerStreamingTypeMapped(Request) returns (stream ResponseTypeMapped);
 
-    rpc ClientStreaming(stream Request) returns (Response);
+  rpc ClientStreaming(stream Request) returns (Response);
 
-    rpc BidiStreaming(stream Request) returns (stream Response);
+  rpc BidiStreaming(stream Request) returns (stream Response);
 }