[LIVY-786] Fix cancellation of driver side jobs

repl/src/main/resources/fake_shell.py

@@ -23,6 +23,7 @@
 import io
 import json
 import logging
+import signal
 import sys
 import traceback
 import base64
@@ -49,6 +50,14 @@
 
 TOP_FRAME_REGEX = re.compile(r'\s*File "<stdin>".*in <module>')
 
+class StatementCancellationException(Exception):
+    pass
+
+def signal_handler(signal, frame):
+    raise StatementCancellationException
+
+signal.signal(signal.SIGUSR1, signal_handler)
+
 def execute_reply(status, content):
     return {
         'msg_type': 'execute_reply',
@@ -648,58 +657,61 @@ def main():
         sys_stdout.flush()
 
         while True:
-            line = sys_stdin.readline()
-
-            if line == '':
-                break
-            elif line == '\n':
-                continue
-
-            try:
-                msg = json.loads(line)
-            except ValueError:
-                LOG.error('failed to parse message', exc_info=True)
-                continue
-
-            try:
-                msg_type = msg['msg_type']
-            except KeyError:
-                LOG.error('missing message type', exc_info=True)
-                continue
-
-            try:
-                content = msg['content']
-            except KeyError:
-                LOG.error('missing content', exc_info=True)
-                continue
-
-            if not isinstance(content, dict):
-                LOG.error('content is not a dictionary')
-                continue
-
             try:
-                handler = msg_type_router[msg_type]
-            except KeyError:
-                LOG.error('unknown message type: %s', msg_type)
+                line = sys_stdin.readline()
+
+                if line == '':
+                    break
+                elif line == '\n':
+                    continue
+
+                try:
+                    msg = json.loads(line)
+                except ValueError:
+                    LOG.error('failed to parse message', exc_info=True)
+                    continue
+
+                try:
+                    msg_type = msg['msg_type']
+                except KeyError:
+                    LOG.error('missing message type', exc_info=True)
+                    continue
+
+                try:
+                    content = msg['content']
+                except KeyError:
+                    LOG.error('missing content', exc_info=True)
+                    continue
+
+                if not isinstance(content, dict):
+                    LOG.error('content is not a dictionary')
+                    continue
+
+                try:
+                    handler = msg_type_router[msg_type]
+                except KeyError:
+                    LOG.error('unknown message type: %s', msg_type)
+                    continue
+
+                response = handler(content)
+
+                try:
+                    response = json.dumps(response)
+                except ValueError:
+                    response = json.dumps({
+                        'msg_type': 'inspect_reply',
+                        'content': {
+                            'status': 'error',
+                            'ename': 'ValueError',
+                            'evalue': 'cannot json-ify %s' % response,
+                            'traceback': [],
+                        }
+                    })
+
+                print(response, file=sys_stdout)
+                sys_stdout.flush()
+            except StatementCancellationException:
                 continue
-
-            response = handler(content)
-
-            try:
-                response = json.dumps(response)
-            except ValueError:
-                response = json.dumps({
-                    'msg_type': 'inspect_reply',
-                    'content': {
-                        'status': 'error',
-                        'ename': 'ValueError',
-                        'evalue': 'cannot json-ify %s' % response,
-                        'traceback': [],
-                    }
-                })
-
-            print(response, file=sys_stdout)
-            sys_stdout.flush()
     finally:
         if os.environ.get("LIVY_TEST") != "true" and 'sc' in global_dict:
             gateway.shutdown_callback_server()

repl/src/main/scala/org/apache/livy/repl/AbstractSparkInterpreter.scala

@@ -61,6 +61,8 @@ abstract class AbstractSparkInterpreter extends Interpreter with Logging {
 
   protected def conf: SparkConf
 
+  private var currentThread: Thread = _
+
   protected def postStart(): Unit = {
     entries = new SparkEntries(conf)
 
@@ -106,10 +108,14 @@ abstract class AbstractSparkInterpreter extends Interpreter with Logging {
   override protected[repl] def execute(code: String): Interpreter.ExecuteResponse =
     restoreContextClassLoader {
       require(isStarted())
-
-      executeLines(code.trim.split("\n").toList, Interpreter.ExecuteSuccess(JObject(
-        (TEXT_PLAIN, JString(""))
-      )))
+      try {
+        currentThread = Thread.currentThread()
+        executeLines(code.trim.split("\n").toList, Interpreter.ExecuteSuccess(JObject(
+          (TEXT_PLAIN, JString(""))
+        )))
+      } finally {
+      currentThread = null
+    }
   }
 
   override protected[repl] def complete(code: String, cursor: Int): Array[String] = {
@@ -349,4 +355,10 @@ abstract class AbstractSparkInterpreter extends Interpreter with Logging {
 
     output
   }
+
+  override def cancel(): Unit = {
+    if (currentThread != null) {
+      currentThread.interrupt()
+    }
+  }
 }

repl/src/main/scala/org/apache/livy/repl/Interpreter.scala

@@ -52,4 +52,7 @@ trait Interpreter {
 
   /** Shut down the interpreter. */
   def close(): Unit
+
+  /** Cancel the executions */
+  def cancel(): Unit
 }

repl/src/main/scala/org/apache/livy/repl/ProcessInterpreter.scala

@@ -76,6 +76,8 @@ abstract class ProcessInterpreter(process: Process)
     }
   }
 
+  override def cancel(): Unit = {}
+
   protected def sendExecuteRequest(request: String): Interpreter.ExecuteResponse
 
   protected def sendShutdownRequest(): Unit = {}

repl/src/main/scala/org/apache/livy/repl/PythonInterpreter.scala

@@ -277,6 +277,20 @@ private class PythonInterpreter(
     }
   }
 
+  override def cancel(): Unit = {
+    val pythonPid = getPythonPid()
+    info("Sending SIGUSR1 to " + pythonPid)
+    Runtime.getRuntime().exec("kill -SIGUSR1 " + pythonPid)
+  }
+
+  def getPythonPid(): Int = {
+    // This implementation is specific to Unix type systems
+    val field = process.getClass().getDeclaredField("pid")
+    field.setAccessible(true)
+    val pid = field.get(process).asInstanceOf[Int]
+    pid
+  }
+
   private def sendRequest(request: Map[String, Any]): Option[JValue] = {
     stdin.println(write(request))
     stdin.flush()

repl/src/main/scala/org/apache/livy/repl/SQLInterpreter.scala

@@ -121,4 +121,6 @@ class SQLInterpreter(
   }
 
   override def close(): Unit = { }
+
+  override def cancel(): Unit = {}
 }

repl/src/main/scala/org/apache/livy/repl/Session.scala

@@ -157,7 +157,7 @@ class Session(
     }
 
     val statementId = newStatementId.getAndIncrement()
-    val statement = new Statement(statementId, code, StatementState.Waiting, null)
+    val statement = new Statement(statementId, code, StatementState.Waiting, null, tpe.name)
     _statements.synchronized { _statements(statementId) = statement }
 
     Future {
@@ -213,6 +213,11 @@ class Session(
           statement.compareAndTransit(StatementState.Cancelling, StatementState.Cancelled)
         } else {
           sc.cancelJobGroup(statementId.toString)
+          val intpOpt = interpreter(Kind(statement.kind))
+          if (!intpOpt.isEmpty) {
+            val intp = intpOpt.get
+            intp.cancel()
+          }
           if (statement.state.get() == StatementState.Cancelling) {
             Thread.sleep(livyConf.getTimeAsMs(RSCConf.Entry.JOB_CANCEL_TRIGGER_INTERVAL))
           }

repl/src/test/scala/org/apache/livy/repl/SparkSessionSpec.scala

@@ -209,7 +209,7 @@ class SparkSessionSpec extends BaseSessionSpec(Spark) {
     eventually(timeout(30 seconds), interval(100 millis)) {
       assert(session.statements(stmtId).state.get() == StatementState.Cancelled)
       session.statements(stmtId).output should include (
-        "Job 0 cancelled part of cancelled job group 0")
+        "java.lang.InterruptedException")
     }
   }
 
@@ -232,7 +232,7 @@ class SparkSessionSpec extends BaseSessionSpec(Spark) {
     eventually(timeout(30 seconds), interval(100 millis)) {
       assert(session.statements(stmtId1).state.get() == StatementState.Cancelled)
       session.statements(stmtId1).output should include (
-        "Job 0 cancelled part of cancelled job group 0")
+        "java.lang.InterruptedException")
     }
   }
 

rsc/src/main/java/org/apache/livy/rsc/driver/Statement.java

@@ -30,6 +30,7 @@ public class Statement {
   public double progress;
   public long started = 0;
   public long completed = 0;
+  public transient String kind;
 
   public Statement(Integer id, String code, StatementState state, String output) {
     this.id = id;
@@ -39,6 +40,11 @@ public Statement(Integer id, String code, StatementState state, String output) {
     this.progress = 0.0;
   }
 
+  public Statement(Integer id, String code, StatementState state, String output, String kind) {
+    this(id, code, state, output);
+    this.kind = kind;
+  }
+
   public Statement() {
     this(null, null, null, null);
   }