From 265f45e81a44fdd7dc6eebc00e58e9558e88cde8 Mon Sep 17 00:00:00 2001
From: Anton Sorokin <anton.sorokin@serokell.io>
Date: Fri, 4 Aug 2023 15:27:12 +0300
Subject: [PATCH] [#22] Threat shutdown and exit requests correctly

Problem:
As in [#22], we want to shut down our LSP server correctly.
That means we should close the pipe and end process after an exit request.
This require proper `shutdown : unit -> bool` function passed
to `Jsonrpc2.Make.run` which returns `true` after
the server received an exit request.

Right now both exit and shutdown requests are setting
LSP server's internal var `_quit` to true. It's very intuitive to
use this var in `shutdown`, but we can't do it, since in this case
the server would stop receiving new messages right after the
shutdown request, despite the fact  according to the LSP specification
it had to wait the exit request.

Solution:
Instead of `_quit : bool` use
```
status : [ `Running | `ReceivedShutdown | `ReceivedExit ]
```
and suggest
```
shutdown () = s#get_status = `ReceivedExit
```
in docs for `Jsonrpc2.Make.run` and in the templete/example.
---
 example/template/main.ml |  5 ++++-
 src/jsonrpc2.ml          |  4 ++++
 src/server.ml            | 13 +++++++++----
 3 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/example/template/main.ml b/example/template/main.ml
index e1ad9188..3e085ccb 100644
--- a/example/template/main.ml
+++ b/example/template/main.ml
@@ -81,7 +81,10 @@ class lsp_server =
 let run () =
   let s = new lsp_server in
   let server = Linol_lwt.Jsonrpc2.create_stdio s in
-  let task = Linol_lwt.Jsonrpc2.run server in
+  let task =
+    let shutdown () = s#get_status = `ReceivedExit in
+    Linol_lwt.Jsonrpc2.run ~shutdown server
+  in
   match Linol_lwt.run task with
   | () -> ()
   | exception e ->
diff --git a/src/jsonrpc2.ml b/src/jsonrpc2.ml
index dc0c52a8..79b973bb 100644
--- a/src/jsonrpc2.ml
+++ b/src/jsonrpc2.ml
@@ -302,6 +302,10 @@ module Make (IO : IO) : S with module IO = IO = struct
       IO.return
       @@ Error (E (ErrorCode.InvalidRequest, "content-type must be 'utf-8'"))
 
+  (** [shutdown ()] is called after processing each request to check if the server
+    could wait for new messages.
+    When launching an LSP server using [Server.Make.server], the
+    natural choice for it is [s#get_status = `ReceivedExit] *)
   let run ?(shutdown = fun _ -> false) (self : t) : unit IO.t =
     let process_msg r =
       let module M = Jsonrpc.Packet in
diff --git a/src/server.ml b/src/server.ml
index eb44de86..1e8ec4bc 100644
--- a/src/server.ml
+++ b/src/server.ml
@@ -157,9 +157,14 @@ module Make (IO : IO) = struct
   class virtual server =
     object (self)
       inherit base_server
-      val mutable _quit = false
+
+      val mutable status : [ `Running | `ReceivedShutdown | `ReceivedExit ] =
+        `Running
+
       val docs : (DocumentUri.t, doc_state) Hashtbl.t = Hashtbl.create 16
-      method! must_quit = _quit
+      method get_status = status
+      (** Check if exit or shutdown request was made by the client.
+        @since NEXT_RELEASE *)
 
       method find_doc (uri : DocumentUri.t) : doc_state option =
         try Some (Hashtbl.find docs uri) with Not_found -> None
@@ -313,7 +318,7 @@ module Make (IO : IO) = struct
           match r with
           | Lsp.Client_request.Shutdown ->
             Log.info (fun k -> k "shutdown");
-            _quit <- true;
+            status <- `ReceivedShutdown;
             IO.return ()
           | Lsp.Client_request.Initialize i ->
             Log.debug (fun k -> k "req: initialize");
@@ -583,7 +588,7 @@ module Make (IO : IO) = struct
             ~old_content:(Lsp.Text_document.text old_doc)
             ~new_content:new_st.content
         | Lsp.Client_notification.Exit ->
-          _quit <- true;
+          status <- `ReceivedExit;
           IO.return ()
         | Lsp.Client_notification.DidSaveTextDocument _
         | Lsp.Client_notification.WillSaveTextDocument _