[LIVY-729] Fix livy recover the killed session

server/src/main/scala/org/apache/livy/server/batch/BatchSession.scala

@@ -188,8 +188,14 @@ class BatchSession(
           info(s"Batch session $id created [appid: ${appId.orNull}, state: ${state.toString}, " +
             s"info: ${appInfo.asJavaMap}]")
         case SparkApp.State.FINISHED => _state = SessionState.Success()
-        case SparkApp.State.KILLED => _state = SessionState.Killed()
-        case SparkApp.State.FAILED => _state = SessionState.Dead()
+        case SparkApp.State.KILLED => {
+          _state = SessionState.Killed()
+          sessionStore.remove(RECOVERY_SESSION_TYPE, id)
+        }
+        case SparkApp.State.FAILED => {
+          _state = SessionState.Dead()
+          sessionStore.remove(RECOVERY_SESSION_TYPE, id)
+        }
         case _ =>
       }
     }

server/src/main/scala/org/apache/livy/server/interactive/InteractiveSession.scala

@@ -577,6 +577,17 @@ class InteractiveSession(
     // Since these 2 transitions are triggered by different threads, there's a race condition.
     // Make sure we won't transit from dead to error state.
     val areSameStates = serverSideState.getClass() == newState.getClass()
+
+    if (!areSameStates) {
+      newState match {
+        case _: SessionState.Killed | _: SessionState.Dead =>
+          sessionStore.remove(RECOVERY_SESSION_TYPE, id)
+        case SessionState.ShuttingDown =>
+          sessionStore.remove(RECOVERY_SESSION_TYPE, id)
+        case _ =>
+      }
+    }
+
     val transitFromInactiveToActive = !serverSideState.isActive && newState.isActive
     if (!areSameStates && !transitFromInactiveToActive) {
       debug(s"$this session state change from ${serverSideState} to $newState")

server/src/main/scala/org/apache/livy/server/recovery/FileSystemStateStore.scala

@@ -125,7 +125,11 @@ class FileSystemStateStore(
   }
 
   override def remove(key: String): Unit = {
-    fileContext.delete(absPath(key), false)
+    try {
+      fileContext.delete(absPath(key), false)
+    } catch {
+      case _: FileNotFoundException => warn(s"Failed to remove non-existed file: ${key}")
+    }
   }
 
   private def absPath(key: String): Path = new Path(fsUri.getPath(), key)

server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperStateStore.scala

@@ -110,7 +110,7 @@ class ZooKeeperStateStore(
     try {
       curatorClient.delete().guaranteed().forPath(prefixKey(key))
     } catch {
-      case _: NoNodeException =>
+      case _: NoNodeException => warn(s"Fail to remove non-existed zookeeper node: ${key}")
     }
   }