[controller] Fix bug where the controller was not exiting

Problem found by Ramkumar, great job!

Indeed I think we never had the right code handling this, as observed
in many instances when the server complained about a forced shutdown.

Likely fixes #230

Co-authored-by: Ramkumar Ramachandra <r@artagnon.com>

CHANGES.md

@@ -16,6 +16,8 @@
  - Support workspaces with multiple roots, this is very useful for
    projects that contain several `_CoqProject` files in different
    directories (@ejgallego, #374)
+ - Fix bug that made the server not exit on `exit` LSP notification
+   (@artagnon, @ejgallego, #375, fixes #230)
 
 # coq-lsp 0.1.5.1: Path
 -----------------------

controller/coq_lsp.ml

@@ -58,11 +58,11 @@ let dispatch_or_resume_check ~ofmt ~state =
 
 (* Wrapper for the top-level call *)
 let dispatch_or_resume_check ~ofmt ~state =
-  try dispatch_or_resume_check ~ofmt ~state with
+  try Some (dispatch_or_resume_check ~ofmt ~state) with
   | U.Type_error (msg, obj) ->
     LIO.trace_object msg obj;
-    state
-  | Lsp_exit -> raise Lsp_exit
+    Some state
+  | Lsp_exit -> None
   | exn ->
     (* Note: We should never arrive here from Coq, as every call to Coq should
        be wrapper in Coq.Protect. So hitting this codepath, is effectively a
@@ -76,13 +76,20 @@ let dispatch_or_resume_check ~ofmt ~state =
     LIO.trace "print_exn [OCaml]" (Printexc.to_string exn);
     LIO.trace "print_exn [Coq  ]" Pp.(string_of_ppcmds CErrors.(iprint iexn));
     LIO.trace "print_bt  [OCaml]" bt;
-    state
+    Some state
+
+let end_lsp_process = ref false
 
 let rec process_queue ofmt ~state =
   if Fleche.Debug.sched_wakeup then
     LIO.trace "<- dequeue" (Format.asprintf "%.2f" (Unix.gettimeofday ()));
-  let state = dispatch_or_resume_check ~ofmt ~state in
-  process_queue ofmt ~state
+  match dispatch_or_resume_check ~ofmt ~state with
+  | None ->
+    end_lsp_process := true;
+    (* XXX: this only works on 4.14 *)
+    (* raise Thread.Exit *)
+    Thread.exit ()
+  | Some state -> process_queue ofmt ~state
 
 let process_input (com : LSP.Message.t) =
   if Fleche.Debug.sched_wakeup then
@@ -128,6 +135,10 @@ let coq_init ~fb_queue ~debug =
   let load_plugin = Coq.Loader.plugin_handler None in
   Coq.Init.(coq_init { fb_handler; debug; load_module; load_plugin })
 
+let exit_message () =
+  let message = "server exiting" in
+  LIO.logMessage ~lvl:1 ~message
+
 let lsp_main bt coqlib vo_load_path ml_include_path =
   (* We output to stdout *)
   let ic = stdin in
@@ -151,9 +162,11 @@ let lsp_main bt coqlib vo_load_path ml_include_path =
 
   (* Read JSON-RPC messages and push them to the queue *)
   let rec read_loop () =
-    let msg = LIO.read_request ic in
-    process_input msg;
-    read_loop ()
+    if !end_lsp_process then ()
+    else
+      let msg = LIO.read_request ic in
+      process_input msg;
+      read_loop ()
   in
 
   (* Input/output will happen now *)
@@ -169,14 +182,12 @@ let lsp_main bt coqlib vo_load_path ml_include_path =
 
     let (_ : Thread.t) = Thread.create (fun () -> process_queue oc ~state) () in
 
-    read_loop ()
+    read_loop ();
+    exit_message ()
   with
-  | (LIO.ReadError "EOF" | Lsp_exit) as exn ->
-    let message =
-      "server exiting"
-      ^ if exn = Lsp_exit then "" else " [uncontrolled LSP shutdown]"
-    in
-    LIO.logMessage ~lvl:1 ~message
+  | LIO.ReadError "EOF" ->
+    (* This is a controlled exit, by EOF *)
+    exit_message ()
   | exn ->
     let bt = Printexc.get_backtrace () in
     let exn, info = Exninfo.capture exn in
@@ -185,7 +196,8 @@ let lsp_main bt coqlib vo_load_path ml_include_path =
     LIO.trace "fatal_error [coq iprint]"
       Pp.(string_of_ppcmds CErrors.(iprint (exn, info)));
     LIO.trace "server crash" (exn_msg ^ bt);
-    LIO.logMessage ~lvl:1 ~message:("server crash: " ^ exn_msg)
+    let message = "[uncontrolled LSP shutdown] server crash\n" ^ exn_msg in
+    LIO.logMessage ~lvl:1 ~message
 
 (* Arguments handling *)
 open Cmdliner