Pull incoming messages in 'StreamIngest' in chunks and use batched request calls

Author: seigert

runtime/src/main/scala/fs2/grpc/client/Fs2ClientCall.scala

@@ -109,7 +109,7 @@ class Fs2ClientCall[F[_], Request, Response] private[client] (
   ): Resource[F, Fs2StreamClientCallListener[F, Response]] = {
     val prefetchN = options.prefetchN.max(1)
     val create = Fs2StreamClientCallListener.create[F, Response](request, signalReadiness, dispatcher, prefetchN)
-    val acquire = start(create, md) <* request(prefetchN)
+    val acquire = start(create, md)
     val release = handleExitCase(cancelSucceed = true)
 
     Resource.makeCase(acquire)(release)

runtime/src/main/scala/fs2/grpc/client/StreamIngest.scala

@@ -24,7 +24,7 @@ package grpc
 package client
 
 import cats.implicits._
-import cats.effect.Concurrent
+import cats.effect.{Concurrent, Ref}
 import cats.effect.std.Queue
 
 private[grpc] trait StreamIngest[F[_], T] {
@@ -39,41 +39,63 @@ private[grpc] object StreamIngest {
       request: Int => F[Unit],
       prefetchN: Int
   ): F[StreamIngest[F, T]] =
-    Queue
-      .unbounded[F, Either[Option[Throwable], T]]
-      .map(q => create[F, T](request, prefetchN, q))
+    (Ref[F].of(0), Queue.unbounded[F, Either[Option[Throwable], T]])
+      .mapN((r, q) => create[F, T](request, prefetchN, r, q))
 
   def create[F[_], T](
       request: Int => F[Unit],
       prefetchN: Int,
+      requested: Ref[F, Int],
       queue: Queue[F, Either[Option[Throwable], T]]
   )(implicit F: Concurrent[F]): StreamIngest[F, T] = new StreamIngest[F, T] {
-
-    val limit: Int =
-      math.max(1, prefetchN)
-
-    val ensureMessages: F[Unit] =
-      queue.size.flatMap(qs => request(1).whenA(qs < limit))
+    private val limit: Int = math.max(1, prefetchN)
+    private def updateRequests: F[Unit] = {
+      queue.size.flatMap { queued =>
+        requested.flatModify { requested =>
+          val total = queued + requested
+          val additional = math.max(0, limit - total)
+
+          (
+            requested + additional,
+            request(additional).whenA(additional > 0)
+          )
+        }
+      }
+    }
 
     def onMessage(msg: T): F[Unit] =
-      queue.offer(msg.asRight) *> ensureMessages
+      queue.offer(msg.asRight) *> requested.update(r => math.max(0, r - 1))
 
     def onClose(error: Option[Throwable]): F[Unit] =
       queue.offer(error.asLeft)
 
     val messages: Stream[F, T] = {
-
-      val run: F[Option[T]] =
-        queue.take.flatMap {
-          case Right(v) => ensureMessages *> v.some.pure[F]
-          case Left(Some(error)) => F.raiseError(error)
-          case Left(None) => none[T].pure[F]
+      type S = Either[Option[Throwable], Chunk[T]]
+
+      def zero: S = Chunk.empty.asRight
+      def loop(state: S): F[Option[(Chunk[T], S)]] =
+        state match {
+          case Left(None) => F.pure(none)
+          case Left(Some(err)) => F.raiseError(err)
+          case Right(acc) =>
+            queue.tryTake.flatMap {
+              case Some(Right(value)) => loop((acc ++ Chunk.singleton(value)).asRight)
+              case Some(Left(err)) =>
+                if (acc.isEmpty) loop(err.asLeft)
+                else F.pure((acc.toIndexedChunk, err.asLeft).some)
+              case None =>
+                val await = if (acc.isEmpty) queue.take.flatMap {
+                  case Right(value) => loop(Chunk.singleton(value).asRight)
+                  case Left(err) => loop(err.asLeft)
+                }
+                else F.pure((acc.toIndexedChunk, zero).some)
+
+                updateRequests *> await
+            }
         }
 
-      Stream.repeatEval(run).unNoneTerminate
-
+      Stream.unfoldChunkEval(zero)(loop)
     }
-
   }
 
 }

runtime/src/test/scala/fs2/grpc/client/StreamIngestSuite.scala

@@ -45,9 +45,9 @@ class StreamIngestSuite extends CatsEffectSuite with CatsEffectFunFixtures {
 
     run(prefetchN = 1, takeN = 1, expectedReq = 1, expectedCount = 1) *>
       run(prefetchN = 2, takeN = 1, expectedReq = 2, expectedCount = 1) *>
-      run(prefetchN = 2, takeN = 2, expectedReq = 3, expectedCount = 2) *>
-      run(prefetchN = 1024, takeN = 1024, expectedReq = 2047, expectedCount = 1024) *>
-      run(prefetchN = 1024, takeN = 1023, expectedReq = 2046, expectedCount = 1023)
+      run(prefetchN = 2, takeN = 2, expectedReq = 2, expectedCount = 2) *>
+      run(prefetchN = 1024, takeN = 1024, expectedReq = 1024, expectedCount = 1024) *>
+      run(prefetchN = 1024, takeN = 1023, expectedReq = 1024, expectedCount = 1023)
 
   }
 

runtime/src/test/scala/fs2/grpc/server/ServerSuite.scala

@@ -360,7 +360,7 @@ class ServerSuite extends Fs2GrpcSuite {
 
     tc.tick()
 
-    assertEquals(dummy.requested, 1)
+    assertEquals(dummy.requested, 2)
 
     listener.onMessage("1")
     tc.tick()