Batching only for Json list, more tests

Author: sebaciv

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

@@ -22,7 +22,7 @@ object RestServlet {
   final val DefaultHandleTimeout = 30.seconds
   final val DefaultMaxPayloadSize = 16 * 1024 * 1024L // 16MB
   final val CookieHeader = "Cookie"
-  final val DefaultStreamingBatchSize = 1
+  final val DefaultStreamingBatchSize = 100
   private final val BufferSize = 8192
 
   /**
@@ -32,7 +32,7 @@ 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]]
+   * @param defaultStreamingBatchSize default batch when streaming [[StreamedBody.JsonList]]
    */
   @explicitGenerics def apply[RestApi: RawRest.AsRawRpc : RestMetadata](
     apiImpl: RestApi,
@@ -130,15 +130,14 @@ class RestServlet(
       case binary: StreamedBody.RawBinary =>
         response.setContentType(binary.contentType)
         binary.content
-          .bufferTumbling(stream.customBatchSize.getOrElse(defaultStreamingBatchSize))
-          .consumeWith(Consumer.foreach { batch =>
-            batch.foreach(e => response.getOutputStream.write(e))
+          .consumeWith(Consumer.foreach { chunk =>
+            response.getOutputStream.write(chunk)
             response.getOutputStream.flush()
           })
       case jsonList: StreamedBody.JsonList =>
         response.setContentType(jsonList.contentType)
         jsonList.elements
-          .bufferTumbling(stream.customBatchSize.getOrElse(defaultStreamingBatchSize))
+          .bufferTumbling(jsonList.customBatchSize.getOrElse(defaultStreamingBatchSize))
           .switchIfEmpty(Observable(Seq.empty))
           .zipWithIndex
           .consumeWith(Consumer.foreach { case (batch, idx) =>

rest/.jvm/src/test/scala/io/udash/rest/ServletBasedRestApiTest.scala

@@ -14,7 +14,7 @@ abstract class ServletBasedRestApiTest extends RestApiTest with UsesHttpServer {
 
   protected def setupServer(server: Server): Unit = {
     val servlet = new RestServlet(serverHandle, serverTimeout, maxPayloadSize)
-    val streamingServlet = new RestServlet(streamingServerHandle, serverTimeout, maxPayloadSize)
+    val streamingServlet = new RestServlet(streamingServerHandle, serverTimeout, maxPayloadSize, defaultStreamingBatchSize = 1)
     val handler = new ServletContextHandler()
     handler.addServlet(new ServletHolder(servlet), "/api/*")
     handler.addServlet(new ServletHolder(streamingServlet), "/stream-api/*")

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

@@ -115,7 +115,6 @@ final class JettyRestClient(
                       code = response.getStatus,
                       headers = parseHeaders(response),
                       body = body,
-                      customBatchSize = Opt.Empty,
                     )
                     callback(Success(restResponse))
                   }
@@ -151,7 +150,6 @@ final class JettyRestClient(
                     code = httpResp.getStatus,
                     headers = parseHeaders(httpResp),
                     body = StreamedBody.fromHttpBody(parseHttpBody(httpResp, this)),
-                    customBatchSize = Opt.Empty,
                   )
                   callback(Success(restResponse))
                 } else {

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

@@ -367,7 +367,15 @@ class addResponseHeader(name: String, value: String) extends ResponseAdjuster wi
   override def adjustResponse(response: StreamedRestResponse): StreamedRestResponse = response.header(name, value)
 }
 
+/**
+ * Annotation which may be applied on REST API methods to change streaming batch size when [[StreamedBody.JsonList]]
+ * bodies are used. It has no effect when applied to other body types or non-streaming methods.
+ */
 class streamingResponseBatchSize(size: Int) extends StreamedResponseAdjuster {
   override def adjustResponse(response: StreamedRestResponse): StreamedRestResponse =
-    response.copy(customBatchSize = Opt.some(size))
+    response.body match {
+      case jsonList: StreamedBody.JsonList =>
+        response.copy(body = jsonList.copy(customBatchSize = Opt.some(size)))
+      case _ => response
+    }
 }

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

@@ -121,7 +121,6 @@ final case class StreamedRestResponse(
   code: Int,
   headers: IMapping[PlainValue],
   body: StreamedBody,
-  customBatchSize: Opt[Int] = Opt.Empty,
 ) extends AbstractRestResponse {
 
   def header(name: String, value: String): StreamedRestResponse =

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

@@ -2,7 +2,7 @@ package io.udash
 package rest.raw
 
 import com.avsystem.commons.annotation.explicitGenerics
-import com.avsystem.commons.misc.ImplicitNotFound
+import com.avsystem.commons.misc.{ImplicitNotFound, Opt}
 import com.avsystem.commons.rpc.{AsRaw, AsRawReal, AsReal}
 import com.avsystem.commons.serialization.GenCodec.ReadFailure
 import monix.reactive.Observable
@@ -49,6 +49,7 @@ object StreamedBody extends StreamedBodyLowPrio {
   final case class JsonList(
     elements: Observable[JsonValue],
     charset: String = HttpBody.Utf8Charset,
+    customBatchSize: Opt[Int] = Opt.Empty,
   ) extends NonEmpty {
     val contentType: String = s"${HttpBody.JsonType};charset=$charset"
 

rest/src/test/scala/io/udash/rest/RestApiTest.scala

@@ -146,14 +146,18 @@ trait RestApiTestScenarios extends RestApiTest {
   }
 }
 
-// TODO streaming MORE tests: cancellation, timeouts, errors, errors after sending a few elements, custom format, slow source observable
+// TODO streaming MORE tests: cancellation
 trait StreamingRestApiTestScenarios extends RestApiTest {
 
   "empty GET stream" in {
     testStream(_.simpleStream(0))
   }
 
-  "trivial GET stream" in {
+  "trivial GET stream - single batch" in {
+    testStream(_.simpleStream(1))
+  }
+
+  "trivial GET stream - multi batch" in {
     testStream(_.simpleStream(5))
   }
 

rest/src/test/scala/io/udash/rest/StreamingRestTestApi.scala

@@ -45,6 +45,7 @@ object CustomStream extends GenCodecRestImplicits {
 }
 
 trait StreamingRestTestApi {
+  @streamingResponseBatchSize(3)
   @GET def simpleStream(size: Int): Observable[Int]
 
   @GET def jsonStream: Observable[RestEntity]

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

@@ -122,7 +122,7 @@ class RawRestTest extends AnyFunSuite with ScalaFutures with Matchers {
 
     val bodyReprTask: Task[String] = resp.body match {
       case StreamedBody.Empty => Task.now("")
-      case StreamedBody.JsonList(elements, charset) =>
+      case StreamedBody.JsonList(elements, charset, _) =>
         elements.toListL.map { list =>
           s" application/json;charset=$charset\n[${list.map(_.value).mkString(",")}]"
         }
@@ -506,7 +506,7 @@ 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) =>
         code shouldBe 200
         val elements = castOrFail[StreamedBody.JsonList](body).elements
         whenReady(elements.toListL.runToFuture) { e =>