update

Author: ArnavBalyan

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

@@ -19,7 +19,7 @@ package org.apache.livy.repl
 
 import java.util.{LinkedHashMap => JLinkedHashMap}
 import java.util.Map.Entry
-import java.util.concurrent.{ConcurrentHashMap, Executors}
+import java.util.concurrent.Executors
 import java.util.concurrent.atomic.AtomicInteger
 
 import scala.collection.JavaConverters._
@@ -63,8 +63,6 @@ class Session(
   private val cancelExecutor = ExecutionContext.fromExecutorService(
     Executors.newSingleThreadExecutor())
 
-  private val statementThreads = new ConcurrentHashMap[Int, Thread]()
-
   private implicit val formats = DefaultFormats
 
   private var _state: SessionState = SessionState.NotStarted
@@ -163,25 +161,18 @@ class Session(
     _statements.synchronized { _statements(statementId) = statement }
 
     Future {
-      val currentThread = Thread.currentThread()
-      statementThreads.put(statementId, currentThread)
-      try {
-        setJobGroup(tpe, statementId)
-        statement.compareAndTransit(StatementState.Waiting, StatementState.Running)
-
-        if (statement.state.get() == StatementState.Running) {
-          statement.started = System.currentTimeMillis()
-          statement.output = executeCode(interpreter(tpe), statementId, code)
-        }
+      setJobGroup(tpe, statementId)
+      statement.compareAndTransit(StatementState.Waiting, StatementState.Running)
 
-        statement.compareAndTransit(StatementState.Running, StatementState.Available)
-        statement.compareAndTransit(StatementState.Cancelling, StatementState.Cancelled)
-        statement.updateProgress(1.0)
-        statement.completed = System.currentTimeMillis()
-      } finally {
-        statementThreads.remove(statementId, currentThread)
-        Thread.interrupted()
+      if (statement.state.get() == StatementState.Running) {
+        statement.started = System.currentTimeMillis()
+        statement.output = executeCode(interpreter(tpe), statementId, code)
       }
+
+      statement.compareAndTransit(StatementState.Running, StatementState.Available)
+      statement.compareAndTransit(StatementState.Cancelling, StatementState.Cancelled)
+      statement.updateProgress(1.0)
+      statement.completed = System.currentTimeMillis()
     }(interpreterExecutor)
 
     statementId
@@ -221,7 +212,6 @@ class Session(
           info(s"Failed to cancel statement $statementId.")
           statement.compareAndTransit(StatementState.Cancelling, StatementState.Cancelled)
         } else {
-          Option(statementThreads.get(statementId)).foreach(_.interrupt())
           sc.cancelJobGroup(statementId.toString)
           if (statement.state.get() == StatementState.Cancelling) {
             Thread.sleep(livyConf.getTimeAsMs(RSCConf.Entry.JOB_CANCEL_TRIGGER_INTERVAL))