Customizable response batch size

Author: sebaciv

rest/.jvm/src/main/scala/io/udash/rest/RestServlet.scala

@@ -22,6 +22,7 @@ object RestServlet {
   final val DefaultHandleTimeout = 30.seconds
   final val DefaultMaxPayloadSize = 16 * 1024 * 1024L // 16MB
   final val CookieHeader = "Cookie"
+  final val DefaultStreamingBatchSize = 1
   private final val BufferSize = 8192
 
   /**
@@ -31,21 +32,29 @@ object RestServlet {
    * @param handleTimeout  maximum time the servlet will wait for results returned by REST API implementation
    * @param maxPayloadSize maximum acceptable incoming payload size, in bytes;
    *                       if exceeded, `413 Payload Too Large` response will be sent back
+   * @param defaultStreamingBatchSize default batch size for [[StreamedRestResponse]]
    */
   @explicitGenerics def apply[RestApi: RawRest.AsRawRpc : RestMetadata](
     apiImpl: RestApi,
     handleTimeout: FiniteDuration = DefaultHandleTimeout,
     maxPayloadSize: Long = DefaultMaxPayloadSize,
+    defaultStreamingBatchSize: Int = DefaultStreamingBatchSize,
   )(implicit
     scheduler: Scheduler
   ): RestServlet =
-    new RestServlet(RawRest.asHandleRequestWithStreaming[RestApi](apiImpl), handleTimeout, maxPayloadSize)
+    new RestServlet(
+      handleRequest = RawRest.asHandleRequestWithStreaming[RestApi](apiImpl),
+      handleTimeout = handleTimeout,
+      maxPayloadSize = maxPayloadSize,
+      defaultStreamingBatchSize = defaultStreamingBatchSize,
+    )
 }
 
 class RestServlet(
   handleRequest: RawRest.HandleRequestWithStreaming,
   handleTimeout: FiniteDuration = DefaultHandleTimeout,
   maxPayloadSize: Long = DefaultMaxPayloadSize,
+  defaultStreamingBatchSize: Int = DefaultStreamingBatchSize,
 )(implicit
   scheduler: Scheduler
 ) extends HttpServlet with LazyLogging {
@@ -120,14 +129,16 @@ class RestServlet(
         Task.eval(writeNonEmptyBody(response, single.body))
       case binary: StreamedBody.RawBinary =>
         response.setContentType(binary.contentType)
-        binary.content.bufferTumbling(stream.batchSize).consumeWith(Consumer.foreach { batch =>
-          batch.foreach(e => response.getOutputStream.write(e))
-          response.getOutputStream.flush()
-        })
+        binary.content
+          .bufferTumbling(stream.customBatchSize.getOrElse(defaultStreamingBatchSize))
+          .consumeWith(Consumer.foreach { batch =>
+            batch.foreach(e => response.getOutputStream.write(e))
+            response.getOutputStream.flush()
+          })
       case jsonList: StreamedBody.JsonList =>
         response.setContentType(jsonList.contentType)
         jsonList.elements
-          .bufferTumbling(stream.batchSize)
+          .bufferTumbling(stream.customBatchSize.getOrElse(defaultStreamingBatchSize))
           .switchIfEmpty(Observable(Seq.empty))
           .zipWithIndex
           .consumeWith(Consumer.foreach { case (batch, idx) =>

rest/jetty/src/main/scala/io/udash/rest/jetty/JettyRestClient.scala

@@ -9,8 +9,8 @@ import io.udash.rest.util.Utils
 import io.udash.utils.URLEncoder
 import monix.eval.Task
 import monix.execution.{Callback, Scheduler}
+import monix.reactive.Observable
 import monix.reactive.OverflowStrategy.Unbounded
-import monix.reactive.{MulticastStrategy, Observable}
 import monix.reactive.subjects.{ConcurrentSubject, PublishToOneSubject}
 import org.eclipse.jetty.client.*
 import org.eclipse.jetty.http.{HttpCookie, HttpHeader, MimeTypes}
@@ -31,7 +31,8 @@ import scala.concurrent.duration.*
  * @param client The configured Jetty `HttpClient` instance.
  * @param defaultMaxResponseLength Default maximum size (in bytes) for buffering non-streamed responses.
  * @param defaultTimeout Default timeout for requests.
- */final class JettyRestClient(
+ */
+final class JettyRestClient(
   client: HttpClient,
   defaultMaxResponseLength: Int = JettyRestClient.DefaultMaxResponseLength,
   defaultTimeout: Duration = JettyRestClient.DefaultTimeout,
@@ -115,7 +116,7 @@ import scala.concurrent.duration.*
                       code = response.getStatus,
                       headers = parseHeaders(response),
                       body = body,
-                      batchSize = 1,
+                      customBatchSize = Opt.Empty,
                     )
                     callback(Success(restResponse))
                   }
@@ -149,7 +150,7 @@ import scala.concurrent.duration.*
                     code = httpResp.getStatus,
                     headers = parseHeaders(httpResp),
                     body = StreamedBody.fromHttpBody(parseHttpBody(httpResp, this)),
-                    batchSize = 1,
+                    customBatchSize = Opt.Empty,
                   )
                   callback(Success(restResponse))
                 } else {
@@ -165,8 +166,8 @@ import scala.concurrent.duration.*
   }
 
   /**
-   * Creates a `RawRest.HandleRequest` which handles standard REST requests by buffering the entire response.
-   * This does *not* support streaming responses.
+   * Creates a [[RawRest.HandleRequest]] which handles standard REST requests by buffering the entire response.
+   * This does <b>not</b> support streaming responses.
    *
    * @param baseUrl The base URL for the REST service.
    * @param customMaxResponseLength Optional override for the maximum response length.

rest/src/main/scala/io/udash/rest/annotations.scala

@@ -1,11 +1,12 @@
 package io.udash
 package rest
 
+import com.avsystem.commons.Opt
 import com.avsystem.commons.annotation.{AnnotationAggregate, defaultsToName}
 import com.avsystem.commons.meta.RealSymAnnotation
-import com.avsystem.commons.rpc._
+import com.avsystem.commons.rpc.*
 import com.avsystem.commons.serialization.optionalParam
-import io.udash.rest.raw._
+import io.udash.rest.raw.*
 
 import scala.annotation.StaticAnnotation
 
@@ -362,6 +363,11 @@ class addRequestHeader(name: String, value: String) extends RequestAdjuster {
   * HTTP header to all outgoing responses generated for invocations of that method on the server side.
   */
 class addResponseHeader(name: String, value: String) extends ResponseAdjuster with StreamedResponseAdjuster {
-  def adjustResponse(response: RestResponse): RestResponse = response.header(name, value)
+  override def adjustResponse(response: RestResponse): RestResponse = response.header(name, value)
   override def adjustResponse(response: StreamedRestResponse): StreamedRestResponse = response.header(name, value)
 }
+
+class streamingResponseBatchSize(size: Int) extends StreamedResponseAdjuster {
+  override def adjustResponse(response: StreamedRestResponse): StreamedRestResponse =
+    response.copy(customBatchSize = Opt.some(size))
+}

rest/src/main/scala/io/udash/rest/raw/RawRest.scala

@@ -142,9 +142,9 @@ trait RawRest {
     RawRest.resolveAndHandle(metadata)(handleResolvedWithStreaming)
 
   /**
-   * Handles a resolved REST call and returns a standard RestResponse.
+   * Handles a resolved REST call and returns a standard [[RestResponse]].
    * This method is maintained for backward compatibility with non-streaming clients.
-   * It delegates to handleResolvedWithStreaming and converts any streaming response to a standard response.
+   * It delegates to [[handleResolvedWithStreaming]] and converts any streaming response to a standard response.
    */
   def handleResolved(request: RestRequest, resolved: ResolvedCall): Task[RestResponse] =
     StreamedRestResponse.fallbackToRestResponse(handleResolvedWithStreaming(request, resolved))
@@ -254,7 +254,7 @@ object RawRest extends RawRpcCompanion[RawRest] {
   @implicitNotFound(InvalidTraitMessage)
   implicit def rawRestAsRawNotFound[T]: ImplicitNotFound[AsRaw[RawRest, T]] = ImplicitNotFound()
 
-  // client side
+  // client side without response streaming support
   def fromHandleRequest[Real: AsRealRpc : RestMetadata](handle: HandleRequest): Real =
     RawRest.asReal(new DefaultRawRest(Nil, RestMetadata[Real], RestParameters.Empty, new RawRest.RestRequestHandler {
       override def handleRequest(request: RestRequest): Task[RestResponse] = handle(request)
@@ -266,7 +266,7 @@ object RawRest extends RawRpcCompanion[RawRest] {
   def fromHandleRequestWithStreaming[Real: AsRealRpc : RestMetadata](handleRequest: RawRest.RestRequestHandler): Real =
     RawRest.asReal(new DefaultRawRest(Nil, RestMetadata[Real], RestParameters.Empty, handleRequest))
 
-  // server side
+  // server side without response streaming support
   def asHandleRequest[Real: AsRawRpc : RestMetadata](real: Real): HandleRequest =
     RawRest.asRaw(real).asHandleRequest(RestMetadata[Real])
 

rest/src/main/scala/io/udash/rest/raw/RestResponse.scala

@@ -113,13 +113,13 @@ trait RestResponseLowPrio { this: RestResponse.type =>
 
 /**
  * Streaming REST response containing a status code, headers, and a streamed body.
- * Unlike standard RestResponse, the body content can be delivered incrementally through a reactive stream.
+ * Unlike standard [[RestResponse]], the body content can be delivered incrementally through a reactive stream.
  */
 final case class StreamedRestResponse(
   code: Int,
   headers: IMapping[PlainValue],
   body: StreamedBody,
-  batchSize: Int,
+  customBatchSize: Opt[Int] = Opt.Empty,
 ) extends AbstractRestResponse {
 
   def header(name: String, value: String): StreamedRestResponse =
@@ -132,7 +132,7 @@ final case class StreamedRestResponse(
 object StreamedRestResponse extends StreamedRestResponseLowPrio {
 
   /**
-   * Converts a StreamedRestResponse to a standard RestResponse by materializing streamed content.
+   * Converts a [[StreamedRestResponse]] to a standard [[RestResponse]] by materializing streamed content.
    * This is useful for compatibility with APIs that don't support streaming.
    */
   def fallbackToRestResponse(response: StreamedRestResponse): Task[RestResponse] = {
@@ -158,7 +158,7 @@ object StreamedRestResponse extends StreamedRestResponseLowPrio {
   }
 
   /**
-   * Converts any AbstractRestResponse to a standard RestResponse by materializing streamed content if necessary.
+   * Converts any [[AbstractRestResponse]] to a standard [[RestResponse]] by materializing streamed content if necessary.
    * This is useful for compatibility with APIs that don't support streaming.
    */
   def fallbackToRestResponse(response: Task[AbstractRestResponse]): Task[RestResponse] =
@@ -168,7 +168,7 @@ object StreamedRestResponse extends StreamedRestResponseLowPrio {
     }
 
   def fromHttpError(error: HttpErrorException): StreamedRestResponse =
-    StreamedRestResponse(error.code, IMapping.empty, StreamedBody.fromHttpBody(error.payload), 1)
+    StreamedRestResponse(error.code, IMapping.empty, StreamedBody.fromHttpBody(error.payload))
 
   class LazyOps(private val resp: () => StreamedRestResponse) extends AnyVal {
     def recoverHttpError: StreamedRestResponse = try resp() catch {

rest/src/main/scala/io/udash/rest/raw/StreamedBody.scala

@@ -21,7 +21,6 @@ sealed trait StreamedBody {
       code = defaultStatus,
       headers = IMapping.empty,
       body = this,
-      batchSize = 1,
     )
 }
 object StreamedBody extends StreamedBodyLowPrio {

rest/src/test/scala/io/udash/rest/raw/RawRestTest.scala

@@ -84,6 +84,7 @@ trait RootApi {
   @POST @CustomBody def echoHeaders(headers: Map[String, String]): Future[WithHeaders[Unit]]
 }
 object RootApi extends DefaultRestApiCompanion[RootApi]
+
 class RawRestTest extends AnyFunSuite with ScalaFutures with Matchers {
   implicit def scheduler: Scheduler = Scheduler.global
 
@@ -206,7 +207,7 @@ class RawRestTest extends AnyFunSuite with ScalaFutures with Matchers {
       serverHandleWithStreaming(request).flatMap {
         case stream: StreamedRestResponse => Task.now(stream)
         case resp: RestResponse =>
-          Task(StreamedRestResponse(resp.code, resp.headers, StreamedBody.fromHttpBody(resp.body), 1))
+          Task(StreamedRestResponse(resp.code, resp.headers, StreamedBody.fromHttpBody(resp.body)))
       }
   })
 
@@ -505,13 +506,12 @@ class RawRestTest extends AnyFunSuite with ScalaFutures with Matchers {
       HttpBody.Empty
     )
     whenReady(serverHandleWithStreaming(request).runToFuture) {
-      case StreamedRestResponse(code, headers, body, batchSize) => {
+      case StreamedRestResponse(code, headers, body, batchSize) =>
         code shouldBe 200
         val elements = castOrFail[StreamedBody.JsonList](body).elements
         whenReady(elements.toListL.runToFuture) { e =>
           e shouldBe List(JsonValue("1"), JsonValue("2"), JsonValue("3"), JsonValue("4"))
         }
-      }
       case _ => fail()
     }
   }