feat: add trace and atomic to trace blocking IOs

linol.opam

@@ -14,6 +14,8 @@ depends: [
   "dune" { >= "2.0" }
   "yojson" { >= "1.6" }
   "logs"
+  "atomic"
+  "trace" { >= "0.4" }
   "lsp" { >= "1.14" & < "1.15" }
   "ocaml" { >= "4.12" }
   "odoc" { with-doc }

src/blocking_IO.ml

@@ -1,3 +1,5 @@
+open Common_
+
 type 'a t = 'a
 type nonrec in_channel = in_channel
 type nonrec out_channel = out_channel
@@ -13,6 +15,7 @@ let stdout = stdout
 
 let default_spawn f =
   let run () =
+    let@ _sp = Trace.with_span ~__FILE__ ~__LINE__ "linol.spawn" in
     try f ()
     with e ->
       Log.err (fun k ->
@@ -22,10 +25,13 @@ let default_spawn f =
   ignore (Thread.create run ())
 
 let catch f g = try f () with e -> g e
+let n_bytes_written = Atomic.make 0
+let n_bytes_read = Atomic.make 0
 
 let rec read ic buf i len =
   if len > 0 then (
     let n = input ic buf i len in
+    ignore (Atomic.fetch_and_add n_bytes_read n : int);
     read ic buf (i + n) (len - n)
   )
 
@@ -33,8 +39,10 @@ let read_line = input_line
 
 let write oc b i len =
   output oc b i len;
+  ignore (Atomic.fetch_and_add n_bytes_written len : int);
   flush oc
 
 let write_string oc s =
   output_string oc s;
+  ignore (Atomic.fetch_and_add n_bytes_written (String.length s) : int);
   flush oc

src/blocking_IO.mli

@@ -6,6 +6,12 @@ include
      and type in_channel = in_channel
      and type out_channel = out_channel
 
+val n_bytes_written : int Atomic.t
+(** @since NEXT_RELEASE *)
+
+val n_bytes_read : int Atomic.t
+(** @since NEXT_RELEASE *)
+
 val default_spawn : (unit -> unit) -> unit
 (** Start a new thread.
     @since NEXT_RELEASE *)

src/common_.ml

@@ -0,0 +1,3 @@
+module Trace = Trace_core
+
+let ( let@ ) = ( @@ )

src/dune

@@ -3,4 +3,4 @@
  (public_name linol)
  (private_modules log)
  (flags :standard -warn-error -a+8)
- (libraries yojson lsp logs threads))
+ (libraries yojson lsp logs atomic threads trace.core))

src/jsonrpc2.ml

@@ -1,3 +1,4 @@
+open Common_
 module J = Yojson.Safe
 module Err = Jsonrpc.Response.Error
 
@@ -155,12 +156,15 @@ module Make (IO : IO) : S with module IO = IO = struct
         send_server_notif self msg)
 
   let handle_notification (self : t) (n : Jsonrpc.Notification.t) : unit IO.t =
+    let@ _sp =
+      Trace.with_span ~__FILE__ ~__LINE__ "linol.handle-notification"
+    in
     match Lsp.Client_notification.of_jsonrpc n with
     | Ok n ->
-      with_error_handler self (fun () ->
+      let@ () = with_error_handler self in
       self.s#on_notification n
         ~notify_back:(send_server_notification self)
-            ~server_request:(server_request self))
+        ~server_request:(server_request self)
     | Error e -> IO.failwith (spf "cannot decode notification: %s" e)
 
   let handle_request (self : t) (r : Jsonrpc.Request.t) : unit IO.t =
@@ -183,7 +187,7 @@ module Make (IO : IO) : S with module IO = IO = struct
       (fun () ->
         match Lsp.Client_request.of_jsonrpc r with
         | Ok (Lsp.Client_request.E r) ->
-          protect ~id (fun () ->
+          let@ () = protect ~id in
           let* reply =
             self.s#on_request r ~id
               ~notify_back:(send_server_notification self)
@@ -191,7 +195,7 @@ module Make (IO : IO) : S with module IO = IO = struct
           in
           let reply_json = Lsp.Client_request.yojson_of_result r reply in
           let response = Jsonrpc.Response.ok id reply_json in
-              send_response self response)
+          send_response self response
         | Error e -> IO.failwith (spf "cannot decode request: %s" e))
       (fun e ->
         let message =