From 2687a2e58284f110774fe0fcb8dc67c3353f855c Mon Sep 17 00:00:00 2001
From: Marc Huisinga <mhuisi@protonmail.com>
Date: Tue, 13 Oct 2020 13:49:46 +0200
Subject: [PATCH] fix: loss of user input when worker event is received

---
 src/Lean/Server/FileWorker.lean |  5 -----
 src/Lean/Server/Watchdog.lean   | 33 ++++++++++++++++++---------------
 2 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index ce174b40e9..d369063394 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -70,14 +70,9 @@ pure $ t.map $ fun error => match error with
 | Except.ok e => e
 | Except.error ioError => Except.error (TaskError.ioError ioError)
 
-partial def crash : Nat → Nat
-| n => crash n+1
-
 private partial def runCore (h : FS.Stream) (uri : DocumentUri) (version : Nat) (contents : FileMap) : Snapshot → IO (Except TaskError ElabTask)
 | parent => do
   result ← compileNextCmd contents.source parent;
-  --let v := crash 1;
-  --IO.eprintln v;
   match result with
   | Sum.inl snap => do
     -- TODO(MH): check for interrupt with increased precision
diff --git a/src/Lean/Server/Watchdog.lean b/src/Lean/Server/Watchdog.lean
index a72bc8a6b5..465ec8ed85 100644
--- a/src/Lean/Server/Watchdog.lean
+++ b/src/Lean/Server/Watchdog.lean
@@ -337,20 +337,20 @@ inductive ServerEvent
 | ClientMsg (msg : JsonRpc.Message)
 | ClientError (e : IO.Error)
 
-partial def mainLoop : Unit → ServerM Unit
-| () => do
+def runClientTask : ServerM (Task ServerEvent) := do
+st ← read;
+clientTask ← liftIO $ IO.asTask $ ServerEvent.ClientMsg <$> readLspMessage st.hIn;
+let clientTask := clientTask.map
+  (fun either => match either with
+  | Except.ok ev   => ev
+  | Except.error e => ServerEvent.ClientError e);
+pure clientTask
+
+partial def mainLoop : Task ServerEvent → ServerM Unit
+| clientTask => do
   st ← read;
   workers ← st.fileWorkersRef.get;
 
-  /- Wait for any of the following events to happen:
-     - client sends us a message
-     - a worker does something -/
-  clientTask ← liftIO $ IO.asTask $ ServerEvent.ClientMsg <$> readLspMessage st.hIn;
-  let clientTask := clientTask.map
-    (fun either => match either with
-    | Except.ok ev   => ev
-    | Except.error e => ServerEvent.ClientError e);
-
   let workerTasks := workers.fold
     (fun acc uri fw =>
       match fw.state with
@@ -367,10 +367,12 @@ partial def mainLoop : Unit → ServerM Unit
       writeLspResponse st.hOut id (Json.null)
     | Message.request id method (some params) => do
       handleRequest id method (toJson params);
-      mainLoop ()
+      newClientTask ← runClientTask;
+      mainLoop newClientTask
     | Message.notification method (some params) => do
       handleNotification method (toJson params);
-      mainLoop ()
+      newClientTask ← runClientTask;
+      mainLoop newClientTask
     | _ => throw (userError "Got invalid JSON-RPC message")
 
   | ServerEvent.ClientError e =>
@@ -381,7 +383,7 @@ partial def mainLoop : Unit → ServerM Unit
     match err with
     | WorkerEvent.crashed e => do
       handleCrash uri #[];
-      mainLoop ()
+      mainLoop clientTask
     -- TODO: this internal error does occur
     | WorkerEvent.terminated => throw (userError "Internal server error: got termination event for worker that should have been removed")
 
@@ -399,7 +401,8 @@ st ← read;
 catch
   (do
     _ ← readLspNotificationAs st.hIn "initialized" InitializedParams;
-    mainLoop ();
+    clientTask ← runClientTask;
+    mainLoop clientTask;
     Message.notification "exit" none ← readLspMessage st.hIn
     | throw (userError "Expected an exit notification");
     pure ())