From 032e3ff8f565b7a0f6a476dbeb4dbb49260066c8 Mon Sep 17 00:00:00 2001
From: Rudi Grinberg <me@rgrinberg.com>
Date: Thu, 23 Jun 2022 10:55:17 -0500
Subject: [PATCH] feature(dune): annotate dune diagnostics with pid

when we are connected to more than one dune instance, we annotate them
with a pid so it's clear to the user

Signed-off-by: Rudi Grinberg <me@rgrinberg.com>

ps-id: 4b487e5f-cd4d-42f2-ae20-b02059f4b5c4
---
 dune-project                         |  2 +-
 ocaml-lsp-server.opam                |  2 +-
 ocaml-lsp-server/src/diagnostics.ml  | 36 ++++++++++++++++++++++++++--
 ocaml-lsp-server/src/diagnostics.mli |  4 +++-
 ocaml-lsp-server/src/dune.ml         |  9 +++----
 5 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/dune-project b/dune-project
index da808999f..9f198d9ac 100644
--- a/dune-project
+++ b/dune-project
@@ -49,7 +49,7 @@ possible and does not make any assumptions about IO.
   yojson
   (re (>= 1.5.0))
   (ppx_yojson_conv_lib (>= "v0.14"))
-  dune-rpc
+  (dune-rpc (>= 3.4.0))
   (chrome-trace (>= 3.3.0))
   dyn
   stdune
diff --git a/ocaml-lsp-server.opam b/ocaml-lsp-server.opam
index ae1426b72..5d7ce84f8 100644
--- a/ocaml-lsp-server.opam
+++ b/ocaml-lsp-server.opam
@@ -23,7 +23,7 @@ depends: [
   "yojson"
   "re" {>= "1.5.0"}
   "ppx_yojson_conv_lib" {>= "v0.14"}
-  "dune-rpc"
+  "dune-rpc" {>= "3.4.0"}
   "chrome-trace" {>= "3.3.0"}
   "dyn"
   "stdune"
diff --git a/ocaml-lsp-server/src/diagnostics.ml b/ocaml-lsp-server/src/diagnostics.ml
index 6004e53e1..5c66caf5c 100644
--- a/ocaml-lsp-server/src/diagnostics.ml
+++ b/ocaml-lsp-server/src/diagnostics.ml
@@ -2,6 +2,8 @@ open Import
 
 let ocamllsp_source = "ocamllsp"
 
+let dune_source = "dune"
+
 module Uri = struct
   include Uri
 
@@ -19,7 +21,29 @@ module Id = struct
   let to_dyn = Dyn.opaque
 end
 
-module Dune = Stdune.Id.Make ()
+module Dune = struct
+  module Id = Stdune.Id.Make ()
+
+  module T = struct
+    type t =
+      { pid : Pid.t
+      ; id : Id.t
+      }
+
+    let compare = Poly.compare
+
+    let equal x y = Ordering.is_eq (compare x y)
+
+    let hash = Poly.hash
+
+    let to_dyn = Dyn.opaque
+  end
+
+  include T
+  module C = Comparable.Make (T)
+
+  let gen pid = { pid; id = Id.gen () }
+end
 
 let equal_message =
   (* because the compiler and merlin wrap messages differently *)
@@ -115,9 +139,17 @@ let send =
             let diagnostics = Table.Multi.find t.merlin uri in
             Table.set pending uri
               (range_map_of_unduplicated_diagnostics diagnostics));
-        Table.iter t.dune ~f:(fun per_dune ->
+        let set_dune_source =
+          let annotate_dune_pid = Table.length t.dune > 1 in
+          if annotate_dune_pid then fun pid (d : Diagnostic.t) ->
+            let source = Some (sprintf "dune (pid=%d)" (Pid.to_int pid)) in
+            { d with source }
+          else fun _pid x -> x
+        in
+        Table.foldi ~init:() t.dune ~f:(fun dune per_dune () ->
             Table.iter per_dune ~f:(fun (uri, diagnostic) ->
                 if Uri_set.mem dirty_uris uri then
+                  let diagnostic = set_dune_source dune.pid diagnostic in
                   add_dune_diagnostic pending uri diagnostic));
         t.dirty_uris <-
           (match which with
diff --git a/ocaml-lsp-server/src/diagnostics.mli b/ocaml-lsp-server/src/diagnostics.mli
index 3c2d4046e..47a358dcd 100644
--- a/ocaml-lsp-server/src/diagnostics.mli
+++ b/ocaml-lsp-server/src/diagnostics.mli
@@ -2,6 +2,8 @@ open Import
 
 val ocamllsp_source : string
 
+val dune_source : string
+
 type t
 
 val create :
@@ -16,7 +18,7 @@ val workspace_root : t -> Uri.t
 module Dune : sig
   type t
 
-  val gen : unit -> t
+  val gen : Pid.t -> t
 end
 
 val set :
diff --git a/ocaml-lsp-server/src/dune.ml b/ocaml-lsp-server/src/dune.ml
index 320596a5f..a95da0b75 100644
--- a/ocaml-lsp-server/src/dune.ml
+++ b/ocaml-lsp-server/src/dune.ml
@@ -224,8 +224,8 @@ end = struct
           let promotions = List.map promotions ~f:For_diff.Diff.of_promotion in
           Some (`Assoc [ For_diff.diagnostic_data promotions ]))
     in
-    Diagnostic.create ?relatedInformation ~range ?severity ~source:"dune"
-      ~message ?data ()
+    Diagnostic.create ?relatedInformation ~range ?severity
+      ~source:Diagnostics.dune_source ~message ?data ()
 
   let progress_loop client progress =
     match Progress.should_report_build_progress progress with
@@ -407,7 +407,7 @@ end = struct
       t.state <- Connected (session, where);
       Fiber.return (Ok ())
 
-  let run ({ config; _ } as t) =
+  let run ({ config; source; _ } as t) =
     let* () = Fiber.return () in
     let session, where =
       match t.state with
@@ -421,7 +421,8 @@ end = struct
       ; finish
       ; promotions = String.Map.empty
       ; client = None
-      ; diagnostics_id = Diagnostics.Dune.gen ()
+      ; diagnostics_id =
+          Diagnostics.Dune.gen (Pid.of_int (Registry.Dune.pid source))
       ; id = Id.gen ()
       }
     in