Merge pull request #375 from ejgallego/fix_server_shutdown

[controller] Fix bug where the controller was not exiting

CHANGES.md

@@ -18,6 +18,8 @@
    directories (@ejgallego, #374)
  - Add VS Code commands to start / stop the server (@ejgallego, #377,
    cc #209)
+ - 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,14 @@ 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 ->
+    (* EJGA: Maybe remove Lsp_exit and have dispatch_or_resume_check return an
+       action? *)
+    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 +79,28 @@ 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
+
+(* Do cleanup here if necessary *)
+let exit_message () =
+  let message = "server exiting" in
+  LIO.logMessage ~lvl:1 ~message
+
+let lsp_cleanup () = exit_message ()
 
 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 ->
+    (* As of now, we exit the whole program here, we could try an experiment to
+       invert the threads, so the I/O routine is a thread and process_queue is
+       the main driver *)
+    lsp_cleanup ();
+    (* We can't use [Thread.exit] here as the main thread will be blocked on
+       I/O *)
+    exit 0
+  | Some state -> process_queue ofmt ~state
 
 let process_input (com : LSP.Message.t) =
   if Fleche.Debug.sched_wakeup then
@@ -128,6 +146,9 @@ 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_notification =
+  Lsp.Base.Message.(Notification { method_ = "exit"; params = [] })
+
 let lsp_main bt coqlib vo_load_path ml_include_path =
   (* We output to stdout *)
   let ic = stdin in
@@ -151,9 +172,13 @@ 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 ()
+    match LIO.read_request ic with
+    | None ->
+      (* EOF, push an exit notication to the queue *)
+      process_input exit_notification
+    | Some msg ->
+      process_input msg;
+      read_loop ()
   in
 
   (* Input/output will happen now *)
@@ -170,22 +195,16 @@ let lsp_main bt coqlib vo_load_path ml_include_path =
     let (_ : Thread.t) = Thread.create (fun () -> process_queue oc ~state) () in
 
     read_loop ()
-  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
-  | exn ->
+  with exn ->
     let bt = Printexc.get_backtrace () in
     let exn, info = Exninfo.capture exn in
     let exn_msg = Printexc.to_string exn in
     LIO.trace "fatal error" (exn_msg ^ bt);
     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

controller/lsp_core.ml

@@ -332,7 +332,7 @@ let version () =
 let rec lsp_init_loop ic ofmt ~cmdline ~debug : (string * Coq.Workspace.t) list
     =
   match LIO.read_request ic with
-  | LSP.Message.Request { method_ = "initialize"; id; params } ->
+  | Some (LSP.Message.Request { method_ = "initialize"; id; params }) ->
     (* At this point logging is allowed per LSP spec *)
     let message =
       Format.asprintf "Initializing coq-lsp server %s" (version ())
@@ -348,13 +348,14 @@ let rec lsp_init_loop ic ofmt ~cmdline ~debug : (string * Coq.Workspace.t) list
     in
     List.iter log_workspace workspaces;
     workspaces
-  | LSP.Message.Request { id; _ } ->
+  | Some (LSP.Message.Request { id; _ }) ->
     (* per spec *)
     LSP.mk_request_error ~id ~code:(-32002) ~message:"server not initialized"
     |> LIO.send_json ofmt;
     lsp_init_loop ic ofmt ~cmdline ~debug
-  | LSP.Message.Notification { method_ = "exit"; params = _ } -> raise Lsp_exit
-  | LSP.Message.Notification _ ->
+  | Some (LSP.Message.Notification { method_ = "exit"; params = _ }) | None ->
+    raise Lsp_exit
+  | Some (LSP.Message.Notification _) ->
     (* We can't log before getting the initialize message *)
     lsp_init_loop ic ofmt ~cmdline ~debug
 

lsp/io.ml

@@ -41,10 +41,10 @@ let read_raw_request ic =
 exception ReadError of string
 
 let read_raw_request ic =
-  try read_raw_request ic with
+  try Some (read_raw_request ic) with
   (* if the end of input is encountered while some more characters are needed to
      read the current conversion specification, or the lsp server closes *)
-  | End_of_file -> raise (ReadError "EOF")
+  | End_of_file -> None
   (* if the input does not match the format. *)
   | Scanf.Scan_failure msg
   (* if a conversion to a number is not possible. *)
@@ -119,10 +119,12 @@ let () = log := trace_object
 
 (** Misc helpers *)
 let rec read_request ic =
-  let com = read_raw_request ic in
-  if Fleche.Debug.read then trace_object "read" com;
-  match Base.Message.from_yojson com with
-  | Ok msg -> msg
-  | Error msg ->
-    trace "read_request" ("error: " ^ msg);
-    read_request ic
+  match read_raw_request ic with
+  | None -> None (* EOF *)
+  | Some com -> (
+    if Fleche.Debug.read then trace_object "read" com;
+    match Base.Message.from_yojson com with
+    | Ok msg -> Some msg
+    | Error msg ->
+      trace "read_request" ("error: " ^ msg);
+      read_request ic)

lsp/io.mli

@@ -18,10 +18,10 @@
 (** JSON-RPC input/output *)
 
 (** Read a JSON-RPC request from channel *)
-val read_raw_request : in_channel -> Yojson.Safe.t
+val read_raw_request : in_channel -> Yojson.Safe.t option
 
-(** *)
-val read_request : in_channel -> Base.Message.t
+(** [None] signals [EOF] *)
+val read_request : in_channel -> Base.Message.t option
 
 exception ReadError of string