Allow overriding the traceln function

lib_eio/core/debug.ml

@@ -0,0 +1,42 @@
+type traceln = {
+  traceln : 'a. ?__POS__:string * int * int * int -> ('a, Format.formatter, unit, unit) format4 -> 'a;
+} [@@unboxed]
+
+let traceln_key : traceln Fiber.key = Fiber.create_key ()
+
+let traceln_mutex = Mutex.create ()
+
+let default_traceln ?__POS__:pos fmt =
+  let k go =
+    let b = Buffer.create 512 in
+    let f = Format.formatter_of_buffer b in
+    go f;
+    Option.iter (fun (file, lnum, _, _) -> Format.fprintf f " [%s:%d]" file lnum) pos;
+    Format.pp_close_box f ();
+    Format.pp_print_flush f ();
+    let msg = Buffer.contents b in
+    Ctf.label msg;
+    let lines = String.split_on_char '\n' msg in
+    Mutex.lock traceln_mutex;
+    Fun.protect ~finally:(fun () -> Mutex.unlock traceln_mutex) @@ fun () ->
+    List.iter (Printf.eprintf "+%s\n") lines;
+    flush stderr
+  in
+  Format.kdprintf k ("@[" ^^ fmt)
+
+let traceln ?__POS__ fmt =
+  let traceln =
+    match Fiber.get traceln_key with
+    | Some { traceln } -> traceln
+    | None
+    | exception Unhandled -> default_traceln
+  in
+  traceln ?__POS__ fmt
+
+type t = <
+  traceln : traceln Fiber.key;
+>
+
+let v = object
+  method traceln = traceln_key
+end

lib_eio/core/eio__core.ml

@@ -8,34 +8,13 @@ module Private = struct
   module Waiters = Waiters
   module Ctf = Ctf
   module Fiber_context = Cancel.Fiber_context
+  module Debug = Debug
 
   module Effects = struct
     type 'a enqueue = 'a Suspend.enqueue
     type _ Effect.t +=
       | Suspend = Suspend.Suspend
       | Fork = Fiber.Fork
       | Get_context = Cancel.Get_context
-      | Trace : (?__POS__:(string * int * int * int) -> ('a, Format.formatter, unit, unit) format4 -> 'a) Effect.t
   end
-
-  let traceln_mutex = Mutex.create ()
-
-  let default_traceln ?__POS__:pos fmt =
-    let k go =
-      let b = Buffer.create 512 in
-      let f = Format.formatter_of_buffer b in
-      go f;
-      Option.iter (fun (file, lnum, _, _) -> Format.fprintf f " [%s:%d]" file lnum) pos;
-      Format.pp_close_box f ();
-      Format.pp_print_flush f ();
-      let msg = Buffer.contents b in
-      Ctf.label msg;
-      let lines = String.split_on_char '\n' msg in
-      Mutex.lock traceln_mutex;
-      Fun.protect ~finally:(fun () -> Mutex.unlock traceln_mutex) @@ fun () ->
-      List.iter (Printf.eprintf "+%s\n") lines;
-      flush stderr
-    in
-    Format.kdprintf k ("@[" ^^ fmt)
-
 end

lib_eio/core/eio__core.mli

@@ -540,12 +540,6 @@ module Private : sig
       (** [perform (Fork new_context f)] creates a new fiber and runs [f] in it, with context [new_context].
           [f] must not raise an exception. See {!Fiber.fork}. *)
 
-      | Trace : (?__POS__:(string * int * int * int) -> ('a, Format.formatter, unit, unit) format4 -> 'a) Effect.t
-      (** [perform Trace fmt] writes trace logging to the configured trace output.
-          It must not switch fibers, as tracing must not affect scheduling.
-          If the system is not ready to receive the trace output,
-          the whole domain must block until it is. *)
-
       | Get_context : Fiber_context.t Effect.t
       (** [perform Get_context] immediately returns the current fiber's context (without switching fibers). *)
   end
@@ -611,18 +605,37 @@ module Private : sig
               and must therefore be holding [mutex]. *)
   end
 
-  val traceln_mutex : Stdlib.Mutex.t
-  (** The mutex used to prevent two domains writing to stderr at once.
+  module Debug : sig
+    val traceln :
+      ?__POS__:string * int * int * int ->
+      ('a, Format.formatter, unit, unit) format4 -> 'a
+    (** Writes trace logging using the current fiber's configured traceln function. *)
+
+    val traceln_mutex : Stdlib.Mutex.t
+    (** The mutex used to prevent two domains writing to stderr at once.
+
+        This might be useful if you want to write to it directly yourself,
+        e.g. for a log reporter. *)
+
+    val default_traceln :
+      ?__POS__:string * int * int * int ->
+      ('a, Format.formatter, unit, unit) format4 -> 'a
+      (** [default_traceln] is a suitable default implementation for {!Eio.Std.traceln}.
 
-      This might be useful if you want to write to it directly yourself,
-      e.g. for a log reporter. *)
+          It writes output to stderr, prefixing each line with a "+".
+          If [__POS__] is given, it also displays the file and line number from that.
+          It uses {!traceln_mutex} so that only one domain's output is written at a time. *)
 
-  val default_traceln :
-    ?__POS__:string * int * int * int ->
-    ('a, Format.formatter, unit, unit) format4 -> 'a
-    (** [default_traceln] is a suitable default implementation for {!Eio.Std.traceln}.
+    type traceln = {
+      traceln : 'a. ?__POS__:string * int * int * int -> ('a, Format.formatter, unit, unit) format4 -> 'a;
+    } [@@unboxed]
+
+    type t = <
+      traceln : traceln Fiber.key;
+    >
+
+    val v : t
+    (** Backends should use this for {!Eio.Stdenv.debug}. *)
+  end
 
-        It writes output to stderr, prefixing each line with a "+".
-        If [__POS__] is given, it also displays the file and line number from that.
-        It uses {!mutex} so that only one domain's output is written at a time. *)
 end

lib_eio/eio.ml

@@ -1,19 +1,16 @@
 include Eio__core
 
-let traceln ?__POS__ fmt =
-  try
-    Effect.perform Private.Effects.Trace ?__POS__ fmt
-  with Unhandled ->
-    Private.default_traceln ?__POS__ fmt
-
 module Fibre = Fiber
 
+module Debug = Private.Debug
+let traceln = Debug.traceln
+
 module Std = struct
   module Promise = Promise
   module Fiber = Fiber
   module Fibre = Fiber
   module Switch = Switch
-  let traceln = traceln
+  let traceln = Debug.traceln
 end
 
 module Semaphore = Semaphore
@@ -41,6 +38,7 @@ module Stdenv = struct
     fs : Fs.dir Path.t;
     cwd : Fs.dir Path.t;
     secure_random : Flow.source;
+    debug : Debug.t;
   >
 
   let stdin  (t : <stdin  : #Flow.source; ..>) = t#stdin
@@ -52,4 +50,5 @@ module Stdenv = struct
   let secure_random (t: <secure_random : #Flow.source; ..>) = t#secure_random
   let fs (t : <fs : #Fs.dir Path.t; ..>) = t#fs
   let cwd (t : <cwd : #Fs.dir Path.t; ..>) = t#cwd
+  let debug (t : <debug : 'a; ..>) = t#debug
 end

lib_eio/eio.mli

@@ -137,6 +137,43 @@ module Fs = Fs
 (** Accessing paths on a file-system. *)
 module Path = Path
 
+(** Control over debugging. *)
+module Debug : sig
+  (** Example:
+      {[
+        open Eio.Std
+
+        let my_traceln = {
+          Eio.Debug.traceln = fun ?__POS__:_ fmt -> Fmt.epr ("[custom-trace] " ^^ fmt ^^ "@.")
+        }
+
+        let () =
+          Eio_main.run @@ fun env ->
+          let debug = Eio.Stdenv.debug env in
+          Fiber.with_binding debug#traceln my_traceln @@ fun () ->
+          traceln "Traced with custom function"
+      ]}
+
+      This will output:
+
+      {[ [custom-trace] Traced with custom function ]}
+  *)
+
+  type traceln = Eio__core.Private.Debug.traceln = {
+    traceln : 'a. ?__POS__:string * int * int * int -> ('a, Format.formatter, unit, unit) format4 -> 'a;
+  } [@@unboxed]
+  (** A function that writes trace logging to some trace output.
+
+      It must not switch fibers, as tracing must not affect scheduling.
+      If the system is not ready to receive the trace output,
+      the whole domain must block until it is. *)
+
+  type t = <
+    traceln : traceln Fiber.key;
+  >
+  (** Fiber keys used to control debugging. Use {!Stdenv.debug} to get this. *)
+end
+
 (** The standard environment of a process. *)
 module Stdenv : sig
   (** All access to the outside world comes from running the event loop,
@@ -162,6 +199,7 @@ module Stdenv : sig
     fs : Fs.dir Path.t;
     cwd : Fs.dir Path.t;
     secure_random : Flow.source;
+    debug : Debug.t;
   >
 
   (** {1 Standard streams}
@@ -218,6 +256,10 @@ module Stdenv : sig
   val secure_random : <secure_random : #Flow.source as 'a; ..> -> 'a
   (** [secure_random t] is a source of random bytes suitable for cryptographic purposes. *)
 
+  (** {1 Debugging} *)
+
+  val debug : <debug : <Debug.t; ..> as 'a; ..> -> 'a
+  (** [debug t] provides privileged controls for debugging. *)
 end
 
 (** {1 Errors and Debugging} *)

lib_eio/mock/backend.ml

@@ -49,9 +49,6 @@ let run main =
           | Eio.Private.Effects.Get_context -> Some (fun k ->
               Effect.Deep.continue k fiber
             )
-          | Eio.Private.Effects.Trace -> Some (fun k ->
-              Effect.Deep.continue k Eio.Private.default_traceln
-            )
           | _ -> None
       }
   in

lib_eio_linux/eio_linux.ml

@@ -1132,6 +1132,7 @@ type stdenv = <
   fs : Eio.Fs.dir Eio.Path.t;
   cwd : Eio.Fs.dir Eio.Path.t;
   secure_random : Eio.Flow.source;
+  debug : Eio.Debug.t;
 >
 
 let domain_mgr ~run_event_loop = object (self)
@@ -1244,6 +1245,7 @@ let stdenv ~run_event_loop =
     method fs = (fs :> Eio.Fs.dir Eio.Path.t)
     method cwd = (cwd :> Eio.Fs.dir Eio.Path.t)
     method secure_random = secure_random
+    method debug = Eio.Private.Debug.v
   end
 
 let pipe sw =
@@ -1371,7 +1373,6 @@ let rec run : type a.
               enqueue_at_head st k ();
               fork ~new_fiber f
             )
-          | Eio.Private.Effects.Trace -> Some (fun k -> continue k Eio.Private.default_traceln)
           | Eio_unix.Private.Await_readable fd -> Some (fun k ->
               match Fiber_context.get_error fiber with
               | Some e -> discontinue k e

lib_eio_linux/eio_linux.mli

@@ -68,6 +68,7 @@ type stdenv = <
   fs : Eio.Fs.dir Eio.Path.t;
   cwd : Eio.Fs.dir Eio.Path.t;
   secure_random : Eio.Flow.source;
+  debug : Eio.Debug.t;
 >
 
 val get_fd : <has_fd; ..> -> FD.t

lib_eio_linux/tests/test.ml

@@ -106,7 +106,7 @@ let test_iovec () =
 
 let () =
   let open Alcotest in
-  run "eioio" [
+  run "eio_linux" [
     "io", [
       test_case "copy"          `Quick test_copy;
       test_case "direct_copy"   `Quick test_direct_copy;

lib_eio_luv/eio_luv.ml

@@ -669,6 +669,7 @@ type stdenv = <
   fs : Eio.Fs.dir Eio.Path.t;
   cwd : Eio.Fs.dir Eio.Path.t;
   secure_random : Eio.Flow.source;
+  debug : Eio.Debug.t;
 >
 
 let domain_mgr ~run_event_loop = object (self)
@@ -850,6 +851,7 @@ let stdenv ~run_event_loop =
     method fs = (fs :> Eio.Fs.dir), "."
     method cwd = (cwd :> Eio.Fs.dir), "."
     method secure_random = secure_random
+    method debug = Eio.Private.Debug.v
   end
 
 let rec wakeup ~async ~io_queued run_q =
@@ -889,8 +891,6 @@ let rec run : type a. (_ -> a) -> a = fun main ->
           Some (fun k -> 
             let k = { Suspended.k; fiber } in
             fn loop fiber (enqueue_thread st k))
-        | Eio.Private.Effects.Trace ->
-          Some (fun k -> continue k Eio.Private.default_traceln)
         | Eio.Private.Effects.Fork (new_fiber, f) ->
           Some (fun k -> 
               let k = { Suspended.k; fiber } in

lib_eio_luv/eio_luv.mli

@@ -121,6 +121,7 @@ type stdenv = <
   fs : Eio.Fs.dir Eio.Path.t;
   cwd : Eio.Fs.dir Eio.Path.t;
   secure_random : Eio.Flow.source;
+  debug : Eio.Debug.t;
 >
 
 val get_fd : <has_fd; ..> -> Low_level.File.t

tests/debug.md

@@ -0,0 +1,28 @@
+# Setting up the environment
+
+```ocaml
+# #require "eio_main";;
+# open Eio.Std;;
+```
+
+## Overriding tracing
+
+```ocaml
+# Eio_main.run @@ fun env ->
+  let debug = Eio.Stdenv.debug env in
+  let my_traceln = {
+    Eio.Debug.traceln = fun ?__POS__:_ fmt -> Fmt.epr ("++" ^^ fmt ^^ "@.")
+  } in
+  Fiber.both
+    (fun () ->
+       Fiber.with_binding debug#traceln my_traceln @@ fun () ->
+       Fiber.both
+         (fun () -> traceln "a")
+         (fun () -> Fiber.yield (); traceln "b")
+     )
+     (fun () -> traceln "c");;
+++a
++c
+++b
+- : unit = ()
+```