Update the codes

Change-Id: I24256600ba0c2c4367ea92f8e9f8952251a93dda

server/src/main/scala/com/cloudera/livy/server/interactive/InteractiveSession.scala

@@ -374,6 +374,8 @@ class InteractiveSession(
   _appId = appIdHint
   sessionStore.save(RECOVERY_SESSION_TYPE, recoveryMetadata)
   heartbeat()
+  // Register this class to RSCClient as a session state listener
+  client.foreach(_.registerStateListener(this))
 
   private val app = mockApp.orElse {
     if (livyConf.isRunningOnYarn()) {
@@ -403,9 +405,6 @@ class InteractiveSession(
     }
     uriFuture onFailure { case e => warn("Fail to get rsc uri", e) }
 
-    // Register this class to RSCClient as a session state listener
-    client.get.registerStateListener(this)
-
     // Send a dummy job that will return once the client is ready to be used, and set the
     // state to "idle" at that point.
     client.get.submit(new PingJob()).addListener(new JobHandle.Listener[Void]() {
@@ -557,7 +556,7 @@ class InteractiveSession(
 
   private def ensureRunning(): Unit = synchronized {
     _state match {
-      case SessionState.Running() | SessionState.Idle() | SessionState.Busy() => Unit
+      case SessionState.Idle() | SessionState.Busy() => Unit
       case _ =>
         throw new IllegalStateException(s"Session is in state ${_state}")
     }