Added support for chunked encoding when job result is a Stream[Byte]

Leon Eller · Leon Eller · commit 9d609cd6984f · 2016-04-12T10:47:30.000+10:00
diff --git a/README.md b/README.md
@@ -620,6 +620,8 @@ serialized properly:
 - Array's
 - Anything that implements Product (Option, case classes) -- they will be serialized as lists
 - Maps and Seqs may contain nested values of any of the above
+- If a job result is of scala's Stream[Byte] type it will be serialised directly as a chunk encoded stream.
+  This is useful if your job result payload is large and may cause a timeout serialising as objects.
 
 If we encounter a data type that is not supported, then the entire result will be serialized to a string.
 
diff --git a/job-server/config/local.conf.template b/job-server/config/local.conf.template
@@ -22,6 +22,9 @@ spark {
     filedao {
       rootdir = /tmp/spark-job-server/filedao/data
     }
+
+    # When using chunked transfer encoding with scala Stream job results, this is the size of each chunk
+    result-chunk-size = 1m
   }
 
   # predefined Spark contexts
diff --git a/job-server/src/spark.jobserver/WebApi.scala b/job-server/src/spark.jobserver/WebApi.scala
@@ -1,38 +1,41 @@
 package spark.jobserver
 
+import java.util.NoSuchElementException
 import java.util.concurrent.TimeUnit
+import javax.net.ssl.SSLContext
 
-import akka.actor.{ ActorSystem, ActorRef }
+import akka.actor.{Actor, ActorLogging, ActorRef, ActorSystem, Props}
 import akka.pattern.ask
 import akka.util.Timeout
-import com.typesafe.config.{ Config, ConfigFactory, ConfigException, ConfigRenderOptions }
-import java.util.NoSuchElementException
-import javax.net.ssl.SSLContext
-import ooyala.common.akka.web.{ WebService, CommonRoutes }
+import com.typesafe.config.{Config, ConfigException, ConfigFactory, ConfigRenderOptions}
+import ooyala.common.akka.web.JsonUtils.AnyJsonFormat
+import ooyala.common.akka.web.{CommonRoutes, WebService}
+import org.apache.shiro.SecurityUtils
+import org.apache.shiro.config.IniSecurityManagerFactory
 import org.joda.time.DateTime
 import org.slf4j.LoggerFactory
-import spark.jobserver.util.SparkJobUtils
-import spark.jobserver.util.SSLContextFactory
-import spark.jobserver.routes.DataRoutes
-import scala.concurrent.{Await, ExecutionContext, Future}
-import scala.util.Try
-import spark.jobserver.io.JobInfo
 import spark.jobserver.auth._
-import spray.http.HttpResponse
-import spray.http.MediaTypes
-import spray.http.StatusCodes
+import spark.jobserver.io.JobInfo
+import spark.jobserver.routes.DataRoutes
+import spark.jobserver.util.{SSLContextFactory, SparkJobUtils}
+import spray.can.Http
+import spray.http._
 import spray.httpx.SprayJsonSupport.sprayJsonMarshaller
+import spray.io.ServerSSLEngineProvider
 import spray.json.DefaultJsonProtocol._
-import spray.routing.{ HttpService, Route, RequestContext }
 import spray.routing.directives.AuthMagnet
-import spray.io.ServerSSLEngineProvider
-import org.apache.shiro.config.IniSecurityManagerFactory
-import org.apache.shiro.mgt.SecurityManager
-import org.apache.shiro.SecurityUtils
+import spray.routing.{HttpService, RequestContext, Route}
+
+import scala.concurrent.{Await, ExecutionContext, Future}
+import scala.util.Try
+
 
 object WebApi {
   val StatusKey = "status"
   val ResultKey = "result"
+  val ResultKeyStartBytes = "{\n".getBytes
+  val ResultKeyEndBytes = "}".getBytes
+  val ResultKeyBytes = ("\"" + ResultKey + "\":").getBytes
 
   def badRequest(ctx: RequestContext, msg: String) {
     ctx.complete(StatusCodes.BadRequest, errMap(msg))
@@ -61,6 +64,14 @@ object WebApi {
     Map(ResultKey -> result)
   }
 
+  def resultToByteIterator(jobReport: Map[String, Any], result: Iterator[_]): Iterator[_] = {
+    ResultKeyStartBytes.toIterator ++
+      (jobReport.map(t => Seq(AnyJsonFormat.write(t._1).toString(),
+                AnyJsonFormat.write(t._2).toString()).mkString(":") ).mkString(",") ++
+        (if(jobReport.nonEmpty) "," else "")).getBytes().toIterator ++
+      ResultKeyBytes.toIterator ++ result ++ ResultKeyEndBytes.toIterator
+  }
+
   def formatException(t: Throwable): Any =
     if (t.getCause != null) {
       Map("message" -> t.getMessage,
@@ -112,6 +123,12 @@ class WebApi(system: ActorSystem,
   val DefaultJobLimit = 50
   val StatusKey = "status"
   val ResultKey = "result"
+  val ResultChunkSize = if (config.hasPath("spark.jobserver.result-chunk-size")) {
+    config.getBytes("spark.jobserver.result-chunk-size").toInt
+  }
+  else {
+    100 * 1024
+  }
 
   val contextTimeout = SparkJobUtils.getContextTimeout(config)
   val bindAddress = config.getString("spark.jobserver.bind-address")
@@ -228,6 +245,7 @@ class WebApi(system: ActorSystem,
    */
   def contextRoutes: Route = pathPrefix("contexts") {
     import ContextSupervisor._
+
     import collection.JavaConverters._
     // user authentication
     authenticate(authenticator) { authInfo =>
@@ -384,7 +402,12 @@ class WebApi(system: ActorSystem,
                 val resultFuture = jobInfo ? GetJobResult(jobId)
                 resultFuture.map {
                   case JobResult(_, result) =>
-                    ctx.complete(jobReport ++ resultToTable(result))
+                    result match {
+                      case s: Stream[_] =>
+                        sendStreamingResponse(ctx, ResultChunkSize,
+                          resultToByteIterator(jobReport, s.toIterator))
+                      case _ => ctx.complete(jobReport ++ resultToTable(result))
+                    }
                   case _ =>
                     ctx.complete(jobReport)
                 }
@@ -466,7 +489,13 @@ class WebApi(system: ActorSystem,
                       JobManagerActor.StartJob(appName, classPath, jobConfig, events))(timeout)
                     respondWithMediaType(MediaTypes.`application/json`) { ctx =>
                       future.map {
-                        case JobResult(_, res)       => ctx.complete(resultToTable(res))
+                        case JobResult(_, res) => {
+                          res match {
+                            case s: Stream[_] => sendStreamingResponse(ctx, ResultChunkSize,
+                              resultToByteIterator(Map.empty, s.toIterator))
+                            case _ => ctx.complete(resultToTable(res))
+                          }
+                        }
                         case JobErroredOut(_, _, ex) => ctx.complete(errMap(ex, "ERROR"))
                         case JobStarted(jobId, context, _) =>
                           jobInfo ! StoreJobConfig(jobId, postedJobConfig)
@@ -504,6 +533,43 @@ class WebApi(system: ActorSystem,
     }
   }
 
+  private def sendStreamingResponse(ctx: RequestContext,
+                                    chunkSize: Int,
+                                    byteIterator: Iterator[_]): Unit = {
+    // simple case class whose instances we use as send confirmation message for streaming chunks
+    case class Ok(remaining: Iterator[_])
+    actorRefFactory.actorOf {
+      Props {
+        new Actor with ActorLogging {
+          // we use the successful sending of a chunk as trigger for sending the next chunk
+          ctx.responder ! ChunkedResponseStart(
+            HttpResponse(entity = HttpEntity(MediaTypes.`application/json`,
+              byteIterator.take(chunkSize).map {
+                case c: Byte => c
+              }.toArray))).withAck(Ok(byteIterator))
+
+          def receive: Receive = {
+            case Ok(remaining) =>
+              val arr = remaining.take(chunkSize).map {
+                case c: Byte => c
+              }.toArray
+              if (arr.nonEmpty) {
+                ctx.responder ! MessageChunk(arr).withAck(Ok(remaining))
+              }
+              else {
+                ctx.responder ! ChunkedMessageEnd
+                context.stop(self)
+              }
+            case ev: Http.ConnectionClosed => {
+              log.warning("Stopping response streaming due to {}", ev)
+              context.stop(self)
+            }
+          }
+        }
+      }
+    }
+  }
+
   override def timeoutRoute: Route =
     complete(500, errMap("Request timed out. Try using the /jobs/<jobID>, /jobs APIs to get status/results"))
 
diff --git a/job-server/test/spark.jobserver/WebApiMainRoutesSpec.scala b/job-server/test/spark.jobserver/WebApiMainRoutesSpec.scala
@@ -1,23 +1,16 @@
 package spark.jobserver
 
-import akka.actor.{Actor, Props}
 import com.typesafe.config.ConfigFactory
-import spark.jobserver.io.{JobInfo, JarInfo}
-import org.joda.time.DateTime
-import org.scalatest.{Matchers, FunSpec, BeforeAndAfterAll}
 import spray.http.StatusCodes._
-import spray.routing.HttpService
-import spray.testkit.ScalatestRouteTest
 
 
 // Tests web response codes and formatting
 // Does NOT test underlying Supervisor / JarManager functionality
 // HttpService trait is needed for the sealRoute() which wraps exception handling
 class WebApiMainRoutesSpec extends WebApiSpec {
-  import scala.collection.JavaConverters._
+  import ooyala.common.akka.web.JsonUtils._
   import spray.httpx.SprayJsonSupport._
   import spray.json.DefaultJsonProtocol._
-  import ooyala.common.akka.web.JsonUtils._
 
   val getJobStatusInfoMap = {
     Map(
@@ -135,6 +128,17 @@ class WebApiMainRoutesSpec extends WebApiSpec {
       }
     }
 
+    it("adhoc job with Stream result of sync route should return 200 and chunked result") {
+      val config2 = "foo.baz = booboo"
+      Post("/jobs?appName=foo.stream&classPath=com.abc.meme&sync=true", config2) ~>
+        sealRoute(routes) ~> check {
+        status should be (OK)
+        responseAs[Map[String, Any]] should be (Map(
+          ResultKey -> "1, 2, 3, 4, 5, 6"
+        ))
+      }
+    }
+
     it("should be able to take a timeout param") {
       val config2 = "foo.baz = booboo"
       Post("/jobs?appName=foo&classPath=com.abc.meme&sync=true&timeout=5", config2) ~>
@@ -269,6 +273,15 @@ class WebApiMainRoutesSpec extends WebApiSpec {
       }
     }
 
+    it("should be able to chunk serialize Stream with different types to JSON") {
+      Get("/jobs/_stream") ~> sealRoute(routes) ~> check {
+        status should be (OK)
+        responseAs[Map[String, Any]] should be (
+          getJobStatusInfoMap ++ Map(ResultKey -> "1, 2, 3, 4, 5, 6, 7")
+        )
+      }
+    }
+
     it("should be able to serialize base types (eg float, numbers) to JSON") {
       Get("/jobs/_num") ~> sealRoute(routes) ~> check {
         status should be (OK)
diff --git a/job-server/test/spark.jobserver/WebApiSpec.scala b/job-server/test/spark.jobserver/WebApiSpec.scala
@@ -4,8 +4,8 @@ import akka.actor.{Actor, Props}
 import com.typesafe.config.ConfigFactory
 import spark.jobserver.io.{JobDAOActor, JobInfo, JarInfo}
 import org.joda.time.DateTime
-import org.scalatest.{Matchers, FunSpec, BeforeAndAfterAll}
-import spray.http.StatusCodes._
+import org.scalatest.{BeforeAndAfterAll, FunSpec, Matchers}
+import spark.jobserver.io.{JarInfo, JobInfo}
 import spray.routing.HttpService
 import spray.testkit.ScalatestRouteTest
 
@@ -75,8 +75,12 @@ with ScalatestRouteTest with HttpService {
         sender ! finishedJobInfo
       case GetJobResult("_seq") =>
         sender ! JobResult("_seq", Seq(1, 2, Map("3" -> "three")))
+      case GetJobResult("_stream") =>
+        sender ! JobResult("_stream", "\"1, 2, 3, 4, 5, 6, 7\"".getBytes().toStream)
       case GetJobStatus("_num") =>
         sender ! finishedJobInfo
+      case GetJobStatus("_stream") =>
+        sender ! finishedJobInfo
       case GetJobResult("_num") =>
         sender ! JobResult("_num", 5000)
       case GetJobStatus("_unk") =>
@@ -117,14 +121,22 @@ with ScalatestRouteTest with HttpService {
       case StartJob("err", _, config, _) =>  sender ! JobErroredOut("foo", dt,
                                                         new RuntimeException("oops",
                                                           new IllegalArgumentException("foo")))
-      case StartJob(_, _, config, events)     =>
+      case StartJob("foo", _, config, events)     =>
         statusActor ! Subscribe("foo", sender, events)
         statusActor ! JobStatusActor.JobInit(JobInfo("foo", "context", null, "", dt, None, None))
         statusActor ! JobStarted("foo", "context1", dt)
         val map = config.entrySet().asScala.map { entry => (entry.getKey -> entry.getValue.unwrapped) }.toMap
         if (events.contains(classOf[JobResult])) sender ! JobResult("foo", map)
         statusActor ! Unsubscribe("foo", sender)
 
+      case StartJob("foo.stream", _, config, events)     =>
+        statusActor ! Subscribe("foo.stream", sender, events)
+        statusActor ! JobStatusActor.JobInit(JobInfo("foo.stream", "context", null, "", dt, None, None))
+        statusActor ! JobStarted("foo.stream", "context1", dt)
+        val result = "\"1, 2, 3, 4, 5, 6\"".getBytes().toStream
+        if (events.contains(classOf[JobResult])) sender ! JobResult("foo.stream", result)
+        statusActor ! Unsubscribe("foo.stream", sender)
+
       case GetJobConfig("badjobid") => sender ! NoSuchJobId
       case GetJobConfig(_)          => sender ! config
     }

Original file line number	Diff line number	Diff line change
`@@ -22,6 +22,9 @@ spark {`
`22`	`22`	`filedao {`
`23`	`23`	`rootdir = /tmp/spark-job-server/filedao/data`
`24`	`24`	`}`
	`25`	`+`
	`26`	`+ # When using chunked transfer encoding with scala Stream job results, this is the size of each chunk`
	`27`	`+ result-chunk-size = 1m`
`25`	`28`	`}`
`26`	`29`
`27`	`30`	`# predefined Spark contexts`