Browse Source

[livy] Clean up session state transitions, expose interruption support

Erick Tryzelaar 10 năm trước cách đây
mục cha
commit
02db858

+ 1 - 1
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/Main.scala

@@ -59,6 +59,6 @@ class ScalatraBootstrap extends LifeCycle {
   }
 
   override def destroy(context: ServletContext): Unit = {
-    sessionManager.close()
+    sessionManager.shutdown()
   }
 }

+ 1 - 0
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/SessionFactory.scala

@@ -2,6 +2,7 @@ package com.cloudera.hue.livy.server
 
 import java.util.UUID
 
+import com.cloudera.hue.livy.server.sessions.{Session, SparkYarnSession, SparkProcessSession}
 import com.cloudera.hue.livy.yarn.Client
 import org.apache.hadoop.yarn.conf.YarnConfiguration
 

+ 8 - 7
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/SessionManager.scala

@@ -1,6 +1,7 @@
 package com.cloudera.hue.livy.server
 
 import com.cloudera.hue.livy.Logging
+import com.cloudera.hue.livy.server.sessions.Session
 
 import scala.collection.concurrent.TrieMap
 import scala.concurrent.duration.Duration
@@ -41,20 +42,20 @@ class SessionManager(factory: SessionFactory) extends Logging {
     })
   }
 
-  def close(): Unit = {
-    Await.result(Future.sequence(sessions.values.map(close)), Duration.Inf)
+  def shutdown(): Unit = {
+    Await.result(Future.sequence(sessions.values.map(delete)), Duration.Inf)
     garbageCollector.shutdown()
   }
 
-  def close(sessionId: String): Future[Unit] = {
+  def delete(sessionId: String): Future[Unit] = {
     sessions.get(sessionId) match {
-      case Some(session) => close(session)
+      case Some(session) => delete(session)
       case None => Future.successful(Unit)
     }
   }
 
-  def close(session: Session): Future[Unit] = {
-    session.close().map { case _ =>
+  def delete(session: Session): Future[Unit] = {
+    session.stop().map { case _ =>
         sessions.remove(session.id)
         Unit
     }
@@ -65,7 +66,7 @@ class SessionManager(factory: SessionFactory) extends Logging {
       System.currentTimeMillis() - session.lastActivity > SessionManager.TIMEOUT
     }
 
-    sessions.values.filter(expired).foreach(close)
+    sessions.values.filter(expired).foreach(delete)
   }
 }
 

+ 34 - 8
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/WebApp.scala

@@ -1,5 +1,6 @@
 package com.cloudera.hue.livy.server
 
+import com.cloudera.hue.livy.server.sessions.{SessionFailedtoStart, Session}
 import com.fasterxml.jackson.core.JsonParseException
 import org.json4s.{DefaultFormats, Formats, MappingException}
 import org.scalatra._
@@ -69,14 +70,6 @@ class WebApp(sessionManager: SessionManager)
     }
   }
 
-  delete("/sessions/:sessionId") {
-    val future = sessionManager.close(params("sessionId"))
-
-    // FIXME: this is silently eating exceptions.
-    //new AsyncResult() { val is = for { _ <- future } yield NoContent }
-    Await.result(future, Duration.Inf)
-  }
-
   post("/sessions/:sessionId/statements") {
     val req = parsedBody.extract[ExecuteStatementRequest]
 
@@ -91,6 +84,39 @@ class WebApp(sessionManager: SessionManager)
     }
   }
 
+  post("/sessions/:sessionId/stop") {
+    sessionManager.get(params("sessionId")) match {
+      case Some(session) =>
+        val future = session.stop()
+
+        // FIXME: this is silently eating exceptions.
+        //new AsyncResult() { val is = for { _ <- future } yield NoContent }
+        Await.result(future, Duration.Inf)
+      case None => NotFound("Session not found")
+    }
+  }
+
+  post("/sessions/:sessionId/interrupt") {
+    sessionManager.get(params("sessionId")) match {
+      case Some(session) =>
+        val future = session.interrupt()
+
+        // FIXME: this is silently eating exceptions.
+        //new AsyncResult() { val is = for { _ <- future } yield NoContent }
+        Await.result(future, Duration.Inf)
+      case None => NotFound("Session not found")
+    }
+  }
+
+  delete("/sessions/:sessionId") {
+    val future = sessionManager.delete(params("sessionId"))
+
+    // FIXME: this is silently eating exceptions.
+    //new AsyncResult() { val is = for { _ <- future } yield NoContent }
+    Await.result(future, Duration.Inf)
+  }
+
+
   val getStatement = get("/sessions/:sessionId/statements/:statementId") {
     sessionManager.get(params("sessionId")) match {
       case Some(session) =>

+ 3 - 3
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/SparkProcessSession.scala → apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/SparkProcessSession.scala

@@ -1,4 +1,4 @@
-package com.cloudera.hue.livy.server
+package com.cloudera.hue.livy.server.sessions
 
 import java.lang.ProcessBuilder.Redirect
 
@@ -64,8 +64,8 @@ object SparkProcessSession extends Logging {
 
 private class SparkProcessSession(id: String, process: Process, port: Int) extends SparkWebSession(id, "localhost", port) {
 
-  override def close(): Future[Unit] = {
-    super.close() andThen { case r =>
+  override def stop(): Future[Unit] = {
+    super.stop() andThen { case r =>
       // Make sure the process is reaped.
       process.waitFor()
 

+ 47 - 32
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/SparkWebSession.scala → apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/SparkWebSession.scala

@@ -1,4 +1,4 @@
-package com.cloudera.hue.livy.server
+package com.cloudera.hue.livy.server.sessions
 
 import com.cloudera.hue.livy._
 import dispatch._
@@ -8,15 +8,13 @@ import org.json4s.{DefaultFormats, Formats}
 import scala.annotation.tailrec
 import scala.concurrent.{Future, _}
 
-abstract class SparkWebSession(val id: String, hostname: String, port: Int)
-  extends Session
-  with Logging {
+abstract class SparkWebSession(val id: String, hostname: String, port: Int) extends Session with Logging {
 
   protected implicit def executor: ExecutionContextExecutor = ExecutionContext.global
   protected implicit def jsonFormats: Formats = DefaultFormats
 
   private[this] var _lastActivity = Long.MaxValue
-  private[this] var _state: State = Running()
+  private[this] var _state: State = Idle()
   private[this] val svc = host(hostname, port)
 
   override def lastActivity: Long = _lastActivity
@@ -24,15 +22,17 @@ abstract class SparkWebSession(val id: String, hostname: String, port: Int)
   override def state: State = _state
 
   override def executeStatement(statement: String): Future[ExecuteResponse] = {
-    ensureRunning {
+    ensureIdle {
+      _state = Busy()
       touchLastActivity()
 
       var req = (svc / "statements").setContentType("application/json", "UTF-8")
       req = req << write(ExecuteRequest(statement))
 
-      for {
-        body <- Http(req OK as.json4s.Json)
-      } yield body.extract[ExecuteResponse]
+      Http(req OK as.json4s.Json).map {
+        transition(Idle())
+        _.extract[ExecuteResponse]
+      }
     }
   }
 
@@ -68,57 +68,72 @@ abstract class SparkWebSession(val id: String, hostname: String, port: Int)
     }
   }
   override def interrupt(): Future[Unit] = {
-    close()
+    stop()
   }
 
-  override def close(): Future[Unit] = {
+  override def stop(): Future[Unit] = {
     synchronized {
       _state match {
-        case Running() =>
-          _state = Stopping()
+        case Idle() =>
+          _state = Busy()
 
           Http(svc.DELETE OK as.String).map { case rep =>
             synchronized {
-              _state = Stopped()
+              _state = Dead()
             }
 
             Unit
           }
-        case Stopping() =>
-          @tailrec
-          def waitForStateChange(state: State): Unit = {
-            if (_state == state) {
-              Thread.sleep(1000)
-              waitForStateChange(state)
-            }
+        case Starting() =>
+          Future {
+            waitForStateChangeFrom(Starting(), { stop() })
           }
-
+        case Busy() =>
           Future {
-            waitForStateChange(Stopping())
-
-            if (_state == Stopped()) {
-              Future.successful(Unit)
-            } else {
-              Future.failed(new IllegalStateException("livy-repl did not stop: %s" format _state))
-            }
+            waitForStateChangeFrom(Busy(), { stop() })
           }
-        case Stopped() =>
+        case Dead() =>
           Future.successful(Unit)
       }
     }
   }
 
+  private def transition(state: State) = synchronized {
+    _state = state
+  }
+
+  @tailrec
+  private def waitForStateChangeFrom[A](state: State, f: => A): A = {
+    if (_state == state) {
+      Thread.sleep(1000)
+      waitForStateChangeFrom(state, f)
+    } else {
+      f
+    }
+  }
+
   private def touchLastActivity() = {
     _lastActivity = System.currentTimeMillis()
   }
 
-  private def ensureRunning[A](f: => A) = {
+  private def ensureIdle[A](f: => A) = {
     synchronized {
-      if (_state == Running()) {
+      if (_state == Idle()) {
         f
       } else {
         throw new IllegalStateException("Session is in state %s" format _state)
       }
     }
   }
+
+  private def ensureRunning[A](f: => A) = {
+    synchronized {
+      _state match {
+        case Idle() | Busy() =>
+          f
+        case _ =>
+          throw new IllegalStateException("Session is in state %s" format _state)
+      }
+    }
+  }
 }

+ 3 - 3
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/SparkYarnSession.scala → apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/SparkYarnSession.scala

@@ -1,4 +1,4 @@
-package com.cloudera.hue.livy.server
+package com.cloudera.hue.livy.server.sessions
 
 import com.cloudera.hue.livy.yarn.{Client, Job}
 import org.apache.hadoop.fs.Path
@@ -42,8 +42,8 @@ object SparkYarnSession {
 private class SparkYarnSession(id: String, job: Job, hostname: String, port: Int)
   extends SparkWebSession(id, hostname, port) {
 
-  override def close(): Future[Unit] = {
-    super.close() andThen { case r =>
+  override def stop(): Future[Unit] = {
+    super.stop() andThen { case r =>
       job.waitForFinish(10000)
       r
     }

+ 8 - 7
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/Session.scala → apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/session.scala

@@ -1,15 +1,10 @@
-package com.cloudera.hue.livy.server
+package com.cloudera.hue.livy.server.sessions
 
 import com.cloudera.hue.livy.ExecuteResponse
 
 import scala.concurrent.Future
 
 trait Session {
-  sealed trait State
-  case class Running() extends State
-  case class Stopping() extends State
-  case class Stopped() extends State
-
   def id: String
 
   def lastActivity: Long
@@ -26,7 +21,13 @@ trait Session {
 
   def interrupt(): Future[Unit]
 
-  def close(): Future[Unit]
+  def stop(): Future[Unit]
 }
 
+sealed trait State
+case class Starting() extends State
+case class Idle() extends State
+case class Busy() extends State
+case class Dead() extends State
+
 class SessionFailedtoStart(msg: String) extends Exception(msg) {}

+ 3 - 0
apps/spark/src/spark/job_server_api.py

@@ -93,3 +93,6 @@ class JobServerApi(object):
 
   def fetch_data(self, session, statement):
     return self._root.get('sessions/%s/statements/%s' % (session, statement))
+
+  def cancel(self, session):
+    return self._root.post('sessions/%s/interrupt' % session)

+ 5 - 1
apps/spark/src/spark/models.py

@@ -307,7 +307,11 @@ class SparkApi():
     }
 
   def cancel(self, notebook, snippet):
-    pass
+    api = get_spark_api(self.user)
+    session = _get_snippet_session(notebook, snippet)
+    response = api.cancel(session['id'])
+
+    return {'status': 'canceled'}
 
   def get_log(self, snippet):
     return 'Not available'