Publisher.toStream can demand for more than one element

reactive-streams/src/main/scala/fs2/interop/reactivestreams/StreamSubscriber.scala

@@ -69,15 +69,19 @@ final class StreamSubscriber[F[_], A](
     sub.onError(t)
   }
 
-  def stream(subscribe: F[Unit]): Stream[F, A] = sub.stream(subscribe)
+  def streamChunk(subscribe: F[Unit]): Stream[F, Chunk[A]] = sub.stream(subscribe)
+  def stream(subscribe: F[Unit]): Stream[F, A] = streamChunk(subscribe).unchunks
 
   private def nonNull[B](b: B): Unit = if (b == null) throw new NullPointerException()
 }
 
 object StreamSubscriber {
 
+  def apply[F[_]: Async, A](bufferSize: Long): F[StreamSubscriber[F, A]] =
+    fsm[F, A](bufferSize).map(new StreamSubscriber(_))
+
   def apply[F[_]: Async, A]: F[StreamSubscriber[F, A]] =
-    fsm[F, A].map(new StreamSubscriber(_))
+    apply(bufferSize = 10L)
 
   @deprecated("Use apply method without dispatcher instead", "3.1.4")
   def apply[F[_]: Async, A](dispatcher: Dispatcher[F]): F[StreamSubscriber[F, A]] =
@@ -102,19 +106,23 @@ object StreamSubscriber {
     def onFinalize: F[Unit]
 
     /** producer for downstream */
-    def dequeue1: F[Either[Throwable, Option[A]]]
+    def dequeue1: F[Either[Throwable, Option[Chunk[A]]]]
 
     /** downstream stream */
-    def stream(subscribe: F[Unit])(implicit ev: ApplicativeError[F, Throwable]): Stream[F, A] =
+    def stream(
+        subscribe: F[Unit]
+    )(implicit ev: ApplicativeError[F, Throwable]): Stream[F, Chunk[A]] =
       Stream.bracket(subscribe)(_ => onFinalize) >> Stream
         .eval(dequeue1)
         .repeat
         .rethrow
         .unNoneTerminate
   }
 
-  private[reactivestreams] def fsm[F[_], A](implicit F: Async[F]): F[FSM[F, A]] = {
-    type Out = Either[Throwable, Option[A]]
+  private[reactivestreams] def fsm[F[_], A](
+      bufferSize: Long
+  )(implicit F: Async[F]): F[FSM[F, A]] = {
+    type Out = Either[Throwable, Option[Chunk[A]]]
 
     sealed trait Input
     case class OnSubscribe(s: Subscription) extends Input
@@ -126,9 +134,10 @@ object StreamSubscriber {
 
     sealed trait State
     case object Uninitialized extends State
-    case class Idle(sub: Subscription) extends State
+    case class Idle(sub: Subscription, buffer: Vector[A]) extends State
     case class RequestBeforeSubscription(req: Out => Unit) extends State
-    case class WaitingOnUpstream(sub: Subscription, elementRequest: Out => Unit) extends State
+    case class WaitingOnUpstream(sub: Subscription, buffer: Vector[A], elementRequest: Out => Unit)
+        extends State
     case object UpstreamCompletion extends State
     case object DownstreamCancellation extends State
     case class UpstreamError(err: Throwable) extends State
@@ -142,8 +151,9 @@ object StreamSubscriber {
     def step(in: Input): State => (State, () => Unit) =
       in match {
         case OnSubscribe(s) => {
-          case RequestBeforeSubscription(req) => WaitingOnUpstream(s, req) -> (() => s.request(1))
-          case Uninitialized                  => Idle(s) -> (() => ())
+          case RequestBeforeSubscription(req) =>
+            WaitingOnUpstream(s, Vector.empty, req) -> (() => s.request(bufferSize))
+          case Uninitialized => Idle(s, Vector.empty) -> (() => ())
           case o =>
             val err = new Error(s"received subscription in invalid state [$o]")
             o -> { () =>
@@ -152,31 +162,44 @@ object StreamSubscriber {
             }
         }
         case OnNext(a) => {
-          case WaitingOnUpstream(s, r) => Idle(s) -> (() => r(a.some.asRight))
-          case DownstreamCancellation  => DownstreamCancellation -> (() => ())
+          case WaitingOnUpstream(s, buffer, r) =>
+            val newBuffer = buffer :+ a
+            if (newBuffer.size == bufferSize) {
+              val chunk = Chunk.vector(newBuffer)
+              Idle(s, Vector.empty) -> (() => r(chunk.some.asRight))
+            } else
+              WaitingOnUpstream(s, newBuffer, r) -> (() => ())
+          case DownstreamCancellation => DownstreamCancellation -> (() => ())
           case o =>
             o -> (() => reportFailure(new Error(s"received record [$a] in invalid state [$o]")))
         }
         case OnComplete => {
-          case WaitingOnUpstream(_, r) => UpstreamCompletion -> (() => r(None.asRight))
-          case _                       => UpstreamCompletion -> (() => ())
+          case WaitingOnUpstream(_, buffer, r) =>
+            if (buffer.nonEmpty) {
+              val chunk = Chunk.vector(buffer)
+              UpstreamCompletion -> (() => r(chunk.some.asRight))
+            } else {
+              UpstreamCompletion -> (() => r(None.asRight))
+            }
+          case _ => UpstreamCompletion -> (() => ())
         }
         case OnError(e) => {
-          case WaitingOnUpstream(_, r) => UpstreamError(e) -> (() => r(e.asLeft))
-          case _                       => UpstreamError(e) -> (() => ())
+          case WaitingOnUpstream(_, _, r) => UpstreamError(e) -> (() => r(e.asLeft))
+          case _                          => UpstreamError(e) -> (() => ())
         }
         case OnFinalize => {
-          case WaitingOnUpstream(sub, r) =>
+          case WaitingOnUpstream(sub, _, r) =>
             DownstreamCancellation -> { () =>
               sub.cancel()
               r(None.asRight)
             }
-          case Idle(sub) => DownstreamCancellation -> (() => sub.cancel())
-          case o         => o -> (() => ())
+          case Idle(sub, _) => DownstreamCancellation -> (() => sub.cancel())
+          case o            => o -> (() => ())
         }
         case OnDequeue(r) => {
-          case Uninitialized          => RequestBeforeSubscription(r) -> (() => ())
-          case Idle(sub)              => WaitingOnUpstream(sub, r) -> (() => sub.request(1))
+          case Uninitialized => RequestBeforeSubscription(r) -> (() => ())
+          case Idle(sub, buffer) =>
+            WaitingOnUpstream(sub, buffer, r) -> (() => sub.request(bufferSize))
           case err @ UpstreamError(e) => err -> (() => r(e.asLeft))
           case UpstreamCompletion     => UpstreamCompletion -> (() => r(None.asRight))
           case o                      => o -> (() => r(new Error(s"received request in invalid state [$o]").asLeft))
@@ -196,8 +219,8 @@ object StreamSubscriber {
         def onError(t: Throwable): Unit = nextState(OnError(t))
         def onComplete(): Unit = nextState(OnComplete)
         def onFinalize: F[Unit] = F.delay(nextState(OnFinalize))
-        def dequeue1: F[Either[Throwable, Option[A]]] =
-          F.async_[Either[Throwable, Option[A]]] { cb =>
+        def dequeue1: F[Either[Throwable, Option[Chunk[A]]]] =
+          F.async_[Either[Throwable, Option[Chunk[A]]]] { cb =>
             nextState(OnDequeue(out => cb(Right(out))))
           }
       }

reactive-streams/src/main/scala/fs2/interop/reactivestreams/package.scala

@@ -45,17 +45,30 @@ import org.reactivestreams._
   */
 package object reactivestreams {
 
+  /** Creates a lazy stream from an `org.reactivestreams.Publisher`.
+    *
+    * The publisher only receives a subscriber when the stream is run.
+    */
+  def fromPublisher[F[_]: Async, A](p: Publisher[A], bufferSize: Long): Stream[F, Chunk[A]] =
+    Stream
+      .eval(StreamSubscriber[F, A](bufferSize))
+      .flatMap(s => s.sub.stream(Sync[F].delay(p.subscribe(s))))
+
   /** Creates a lazy stream from an `org.reactivestreams.Publisher`.
     *
     * The publisher only receives a subscriber when the stream is run.
     */
   def fromPublisher[F[_]: Async, A](p: Publisher[A]): Stream[F, A] =
     Stream
       .eval(StreamSubscriber[F, A])
-      .flatMap(s => s.sub.stream(Sync[F].delay(p.subscribe(s))))
+      .flatMap(s => s.sub.stream(Sync[F].delay(p.subscribe(s))).unchunks)
 
   implicit final class PublisherOps[A](val publisher: Publisher[A]) extends AnyVal {
 
+    /** Creates a lazy stream from an `org.reactivestreams.Publisher` */
+    def toStreamChunk[F[_]: Async](bufferSize: Long): Stream[F, Chunk[A]] =
+      fromPublisher(publisher, bufferSize)
+
     /** Creates a lazy stream from an `org.reactivestreams.Publisher` */
     def toStream[F[_]: Async]: Stream[F, A] =
       fromPublisher(publisher)