Merge pull request spark-jobserver#418 from leone17/streaming

velvia · velvia · commit 2c558af193f6 · 2016-04-13T18:01:49.000-06:00
Added support for chunked encoding when job result is a Stream[Byte]
diff --git a/README.md b/README.md
@@ -629,6 +629,11 @@ 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. Beware, this
+  will not currently work as desired with context-per-jvm=true configuration, since it would require serialising
+  Stream[_] blob between processes. For now use Stream[_] job results in context-per-jvm=false configuration, pending
+  potential future enhancements to support this in context-per-jvm=true mode.
 
 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/ChunkEncodedStreamingSupport.scala b/job-server/src/spark.jobserver/ChunkEncodedStreamingSupport.scala
@@ -0,0 +1,18 @@
+package spark.jobserver
+
+import akka.actor.Props
+import spray.routing.{HttpService, RequestContext}
+
+trait ChunkEncodedStreamingSupport {
+  this: HttpService =>
+
+  protected def sendStreamingResponse(ctx: RequestContext,
+                                    chunkSize: Int,
+                                    byteIterator: Iterator[_]): Unit = {
+    actorRefFactory.actorOf {
+      Props {
+        new ChunkEncodingActor(ctx, chunkSize, byteIterator)
+      }
+    }
+  }
+}
diff --git a/job-server/src/spark.jobserver/ChunkEncodingActor.scala b/job-server/src/spark.jobserver/ChunkEncodingActor.scala
@@ -0,0 +1,48 @@
+package spark.jobserver
+
+import akka.actor.{Actor, ActorLogging}
+import spark.jobserver.ChunkEncodingActor.Ok
+import spray.can.Http
+import spray.http._
+import spray.routing.RequestContext
+
+object ChunkEncodingActor {
+  // simple case class whose instances we use as send confirmation message for streaming chunks
+  case class Ok(remaining: Iterator[_])
+}
+
+/**
+  * Performs sending back a response in streaming fashion using chunk encoding
+  * @param ctx RequestContext which has the responder to send chunks to
+  * @param chunkSize The size of each chunk
+  * @param byteIterator Iterator of data to stream back (currently only Byte is supported)
+  */
+class ChunkEncodingActor(ctx: RequestContext,
+                         chunkSize: Int,
+                         byteIterator: Iterator[_]) extends 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)
+    }
+  }
+}
+
diff --git a/job-server/src/spark.jobserver/JobManagerActor.scala b/job-server/src/spark.jobserver/JobManagerActor.scala
@@ -307,6 +307,15 @@ class JobManagerActor(contextConfig: Config) extends InstrumentedActor {
     }(executionContext).andThen {
       case Success(result: Any) =>
         statusActor ! JobFinished(jobId, DateTime.now())
+        // TODO: If the result is Stream[_] and this is running with context-per-jvm=true configuration
+        // serializing a Stream[_] blob across process boundaries is not desirable.
+        // In that scenario an enhancement is required here to chunk stream results back.
+        // Something like ChunkedJobResultStart, ChunkJobResultMessage, and ChunkJobResultEnd messages
+        // might be a better way to send results back and then on the other side use chunked encoding
+        // transfer to send the chunks back. Alternatively the stream could be persisted here to HDFS
+        // and the streamed out of InputStream on the other side.
+        // Either way an enhancement would be required here to make Stream[_] responses work
+        // with context-per-jvm=true configuration
         resultActor ! JobResult(jobId, result)
       case Failure(error: Throwable) =>
         // Wrapping the error inside a RuntimeException to handle the case of throwing custom exceptions.
diff --git a/job-server/src/spark.jobserver/WebApi.scala b/job-server/src/spark.jobserver/WebApi.scala
@@ -1,38 +1,40 @@
 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.{ActorRef, ActorSystem}
 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.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 +63,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,
@@ -95,7 +105,8 @@ class WebApi(system: ActorSystem,
              dataManager: ActorRef,
              supervisor: ActorRef,
              jobInfo: ActorRef)
-    extends HttpService with CommonRoutes with DataRoutes with SJSAuthenticator with CORSSupport {
+    extends HttpService with CommonRoutes with DataRoutes with SJSAuthenticator with CORSSupport
+                        with ChunkEncodedStreamingSupport {
   import CommonMessages._
   import ContextSupervisor._
   import scala.concurrent.duration._
@@ -112,6 +123,8 @@ class WebApi(system: ActorSystem,
   val DefaultJobLimit = 50
   val StatusKey = "status"
   val ResultKey = "result"
+  val ResultChunkSize = Option("spark.jobserver.result-chunk-size").filter(config.hasPath)
+      .fold(100 * 1024)(config.getBytes(_).toInt)
 
   val contextTimeout = SparkJobUtils.getContextTimeout(config)
   val bindAddress = config.getString("spark.jobserver.bind-address")
@@ -228,6 +241,7 @@ class WebApi(system: ActorSystem,
    */
   def contextRoutes: Route = pathPrefix("contexts") {
     import ContextSupervisor._
+
     import collection.JavaConverters._
     // user authentication
     authenticate(authenticator) { authInfo =>
@@ -384,7 +398,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 +485,12 @@ 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)
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`