Improve tracing

This adds tracing of cancellation contexts and OS operations, and
documents and cleans up the API a bit.

When reading events, instead of the consumer providing one function per
event, we now use a single function taking a variant. This makes the API
easier to use and lets the caller decide whether they want to handle all
events (checked by the compiler) or just some of them.

Note: eio_posix yields before each operation to avoid starvation.
However, if it needed to block then it did another yield afterwards when
retrying, which isn't necessary and clutters up the traces. This commit
also removes that.

Co-authored-by: Lucas Pluvinage <[email protected]>

Author: talex5

examples/trace/main.ml

@@ -9,13 +9,6 @@
 
 open Eio.Std
 
-(* [handle pp] handles an event by writing the timestamp, ring ID and user data with [traceln].
-   [pp] is used to format the user data. *)
-let handle pp : _ Eio_runtime_events.handler =
-  fun ring ts v ->
-  (* Note: don't use traceln here, as it will just generate more log events! *)
-  Fmt.epr "%9Ld:ring %d: %a@." (Runtime_events.Timestamp.to_int64 ts) ring pp v
-
 let callbacks =
   Runtime_events.Callbacks.create ()
     (* Uncomment to trace GC events too: *)
@@ -25,11 +18,10 @@ let callbacks =
 *)
     ~lost_events:(fun ring n -> traceln "ring %d lost %d events" ring n)
   |> Eio_runtime_events.add_callbacks
-    ~fiber:(handle (Fmt.fmt "running fiber %d"))
-    ~create:(handle (fun f (id, ty) -> Fmt.pf f "create %s %d" (Eio_runtime_events.ty_to_string ty) id))
-    ~resolve:(handle (Fmt.fmt "resolve %d"))
-    ~log:(handle (Fmt.fmt "log %S"))
-    ~name:(handle (fun f (id, name) -> Fmt.pf f "%d is named %S" id name))
+    (fun ring ts e ->
+       (* Note: don't use traceln here, as it will just generate more log events! *)
+       Fmt.epr "%9Ld:ring %d: %a@." (Runtime_events.Timestamp.to_int64 ts) ring Eio_runtime_events.pp_event e
+    )
     (* (see lib_eio/runtime_events/eio_runtime_events.mli for more event types) *)
 
 (* Read and display trace events from [cursor] until [finished]. *)

lib_eio/buf_write.ml

@@ -456,7 +456,7 @@ let rec await_batch t =
   | Closed, false -> raise End_of_file
   | (Active | Closed), true -> Buffers.to_list t.scheduled
   | Paused, _ | Active, false ->
-    Suspend.enter (fun ctx enqueue ->
+    Suspend.enter "Buf_write.await_batch" (fun ctx enqueue ->
         Fiber_context.set_cancel_fn ctx (fun ex ->
             t.wake_writer <- ignore;
             enqueue (Error ex)

lib_eio/condition.ml

@@ -12,7 +12,7 @@ let lock_protected m =
 
 let await_generic ?mutex t =
   match
-    Suspend.enter_unchecked (fun ctx enqueue ->
+    Suspend.enter_unchecked "Condition.await" (fun ctx enqueue ->
         match Fiber_context.get_error ctx with
         | Some ex ->
           Option.iter Eio_mutex.unlock mutex;
@@ -90,7 +90,7 @@ let rec loop_no_mutex t fn =
     ensure_cancelled request;
     x
   | None ->
-    Suspend.enter_unchecked (fun ctx enqueue ->
+    Suspend.enter_unchecked "Condition.loop_no_mutex" (fun ctx enqueue ->
         match Fiber_context.get_error ctx with
         | Some ex ->
           ensure_cancelled request;

lib_eio/core/cancel.ml

@@ -14,6 +14,7 @@ type state =
    If a function can succeed in a separate domain,
    the user's cancel function is responsible for ensuring that this is done atomically. *)
 type t = {
+  id : Trace.id;
   mutable state : state;
   children : t Lwt_dllist.t;
   fibers : fiber_context Lwt_dllist.t;
@@ -89,10 +90,12 @@ let move_fiber_to t fiber =
   fiber.cancel_node <- Some new_node
 
 (* Note: the new value is not linked into the cancellation tree. *)
-let create ~protected =
+let create ~protected purpose =
   let children = Lwt_dllist.create () in
   let fibers = Lwt_dllist.create () in
-  { state = Finished; children; protected; fibers; domain = Domain.self () }
+  let id = Trace.mint_id () in
+  Trace.create_cc id purpose;
+  { id; state = Finished; children; protected; fibers; domain = Domain.self () }
 
 (* Links [t] into the tree as a child of [parent] and returns a function to remove it again. *)
 let activate t ~parent =
@@ -106,19 +109,19 @@ let activate t ~parent =
     Lwt_dllist.remove node
 
 (* Runs [fn] with a fresh cancellation context. *)
-let with_cc ~ctx:fiber ~parent ~protected fn =
+let with_cc ~ctx:fiber ~parent ~protected purpose fn =
   if not protected then check parent;
-  let t = create ~protected in
+  let t = create ~protected purpose in
   let deactivate = activate t ~parent in
   move_fiber_to t fiber;
   let cleanup () = move_fiber_to parent fiber; deactivate () in
   match fn t with
-  | x            -> cleanup (); x
-  | exception ex -> cleanup (); raise ex
+  | x            -> cleanup (); Trace.exit_cc (); x
+  | exception ex -> cleanup (); Trace.error t.id ex; Trace.exit_cc (); raise ex
 
 let protect fn =
   let ctx = Effect.perform Get_context in
-  with_cc ~ctx ~parent:ctx.cancel_context ~protected:true @@ fun _ ->
+  with_cc ~ctx ~parent:ctx.cancel_context ~protected:true Protect @@ fun _ ->
   (* Note: there is no need to check the new context after [fn] returns;
      the goal of cancellation is only to finish the thread promptly, not to report the error.
      We also do not check the parent context, to make sure the caller has a chance to handle the result. *)
@@ -167,18 +170,21 @@ let cancel t ex =
       Printexc.raise_with_backtrace ex bt
   )
 
-let sub fn =
+let sub_checked purpose fn =
   let ctx = Effect.perform Get_context in
   let parent = ctx.cancel_context in
-  with_cc ~ctx ~parent ~protected:false @@ fun t ->
+  with_cc ~ctx ~parent ~protected:false purpose @@ fun t ->
   fn t
 
+let sub fn =
+  sub_checked Sub fn
+
 (* Like [sub], but it's OK if the new context is cancelled.
    (instead, return the parent context on exit so the caller can check that) *)
-let sub_unchecked fn =
+let sub_unchecked purpose fn =
   let ctx = Effect.perform Get_context in
   let parent = ctx.cancel_context in
-  with_cc ~ctx ~parent ~protected:false @@ fun t ->
+  with_cc ~ctx ~parent ~protected:false purpose @@ fun t ->
   fn t;
   parent
 
@@ -198,17 +204,18 @@ module Fiber_context = struct
 
   let make ~cc ~vars =
     let tid = Trace.mint_id () in
-    Trace.create tid Fiber;
+    Trace.create_fiber tid ~cc:cc.id;
     let t = { tid; cancel_context = cc; cancel_node = None; cancel_fn = ignore; vars } in
     t.cancel_node <- Some (Lwt_dllist.add_r t cc.fibers);
     t
 
   let make_root () =
-    let cc = create ~protected:false in
+    let cc = create ~protected:false Root in
     cc.state <- On;
     make ~cc ~vars:Hmap.empty
 
   let destroy t =
+    Trace.exit_fiber t.tid;
     Option.iter Lwt_dllist.remove t.cancel_node
 
   let vars t = t.vars

lib_eio/core/dla.ml

@@ -7,7 +7,7 @@ let prepare_for_await () =
       | _ -> ()
   and await () =
     if Atomic.get state != `Released then
-      Suspend.enter @@ fun ctx enqueue ->
+      Suspend.enter "domain-local-await" @@ fun ctx enqueue ->
       let awaiting = `Awaiting enqueue in
       if Atomic.compare_and_set state `Init awaiting then (
         Cancel.Fiber_context.set_cancel_fn ctx (fun ex ->

lib_eio/core/eio__core.mli

@@ -698,15 +698,18 @@ module Private : sig
 
   (** Suspend a fiber and enter the scheduler. *)
   module Suspend : sig
-    val enter : (Fiber_context.t -> 'a Effects.enqueue -> unit) -> 'a
-    (** [enter fn] suspends the calling fiber and calls [fn ctx enqueue] in the scheduler's context.
+    val enter : string -> (Fiber_context.t -> 'a Effects.enqueue -> unit) -> 'a
+    (** [enter op fn] suspends the calling fiber and calls [fn ctx enqueue] in the scheduler's context.
+
         This should arrange for [enqueue] to be called when the fiber should be resumed.
         [enqueue] is thread-safe and so can be called from another domain or systhread.
 
         [ctx] should be used to set a cancellation function. Otherwise, the operation is non-interruptable.
-        If the caller's cancellation context is already cancelled, [enter] immediately aborts. *)
+        If the caller's cancellation context is already cancelled, [enter] immediately aborts.
+
+        [op] is used when tracing to label the operation. *)
 
-    val enter_unchecked : (Fiber_context.t -> 'a Effects.enqueue -> unit) -> 'a
+    val enter_unchecked : string -> (Fiber_context.t -> 'a Effects.enqueue -> unit) -> 'a
     (** [enter_unchecked] is like [enter] except that it does not perform the initial check
         that the fiber isn't cancelled (this is useful if you want to do the check yourself, e.g.
         because you need to unlock a mutex if cancelled). *)

lib_eio/core/fiber.ml

@@ -3,7 +3,7 @@
 type _ Effect.t += Fork : Cancel.fiber_context * (unit -> unit) -> unit Effect.t
 
 let yield () =
-  let fiber = Suspend.enter (fun fiber enqueue -> enqueue (Ok fiber)) in
+  let fiber = Suspend.enter "" (fun fiber enqueue -> enqueue (Ok fiber)) in
   Cancel.check fiber.cancel_context
 
 (* Note: [f] must not raise an exception, as that would terminate the whole scheduler. *)
@@ -17,12 +17,11 @@ let fork ~sw f =
     let new_fiber = Cancel.Fiber_context.make ~cc:sw.cancel ~vars in
     fork_raw new_fiber @@ fun () ->
     Switch.with_op sw @@ fun () ->
-    match f () with
-    | () ->
-      Trace.resolve (Cancel.Fiber_context.tid new_fiber)
-    | exception ex ->
-      Switch.fail sw ex;  (* The [with_op] ensures this will succeed *)
-      Trace.resolve_error (Cancel.Fiber_context.tid new_fiber) ex
+    try
+      f ()
+    with ex ->
+      let bt = Printexc.get_raw_backtrace () in
+      Switch.fail ~bt sw ex;  (* The [with_op] ensures this will succeed *)
   ) (* else the fiber should report the error to [sw], but [sw] is failed anyway *)
 
 let fork_daemon ~sw f =
@@ -35,13 +34,12 @@ let fork_daemon ~sw f =
     match f () with
     | `Stop_daemon ->
       (* The daemon asked to stop. *)
-      Trace.resolve (Cancel.Fiber_context.tid new_fiber)
+      ()
     | exception Cancel.Cancelled Exit when not (Cancel.is_on sw.cancel) ->
       (* The daemon was cancelled because all non-daemon fibers are finished. *)
-      Trace.resolve (Cancel.Fiber_context.tid new_fiber)
+      ()
     | exception ex ->
       Switch.fail sw ex;  (* The [with_daemon] ensures this will succeed *)
-      Trace.resolve_error (Cancel.Fiber_context.tid new_fiber) ex
   ) (* else the fiber should report the error to [sw], but [sw] is failed anyway *)
 
 let fork_promise ~sw f =
@@ -86,13 +84,13 @@ let pair f g =
 exception Not_first
 
 let await_cancel () =
-  Suspend.enter @@ fun fiber enqueue ->
+  Suspend.enter "await_cancel" @@ fun fiber enqueue ->
   Cancel.Fiber_context.set_cancel_fn fiber (fun ex -> enqueue (Error ex))
 
 let any fs =
   let r = ref `None in
   let parent_c =
-    Cancel.sub_unchecked (fun cc ->
+    Cancel.sub_unchecked Any (fun cc ->
         let wrap h =
           match h () with
           | x ->
@@ -198,7 +196,7 @@ module List = struct
       }
 
     let await_free t =
-      if t.free_fibers = 0 then Single_waiter.await t.cond t.sw.id;
+      if t.free_fibers = 0 then Single_waiter.await t.cond "Limiter.await_free" t.sw.cancel.id;
       (* If we got woken up then there was a free fiber then. And since we're the
          only fiber that uses [t], and we were sleeping, it must still be free. *)
       assert (t.free_fibers > 0);
@@ -306,7 +304,7 @@ let unwrap_cancelled state =
 let run_coroutine ~state fn =
   let await_request ~prev ~on_suspend =
     (* Suspend and wait for the consumer to resume us: *)
-    Suspend.enter (fun ctx enqueue ->
+    Suspend.enter "await-consumer" (fun ctx enqueue ->
         let ready = `Ready enqueue in
         if Atomic.compare_and_set state prev ready then (
           Cancel.Fiber_context.set_cancel_fn ctx (fun ex ->
@@ -355,7 +353,7 @@ let fork_coroutine ~sw fn =
             raise ex
     );
   fun () ->
-    Suspend.enter (fun ctx enqueue ->
+    Suspend.enter "await-producer" (fun ctx enqueue ->
         let rec aux () =
           match Atomic.get state with
           | `Ready resume as prev ->

lib_eio/core/promise.ml

@@ -26,22 +26,22 @@ let create_with_id id =
 
 let create ?label () =
   let id = Trace.mint_id () in
-  Trace.create ?label id Promise;
+  Trace.create_obj ?label id Promise;
   create_with_id id
 
 let create_resolved x =
   let id = Trace.mint_id () in
-  Trace.create id Promise;
+  Trace.create_obj id Promise;
   to_public_promise { id; state = Atomic.make (Resolved x) }
 
 let await t =
   let t = of_public_promise t in
   match Atomic.get t.state with
   | Resolved x ->
-    Trace.read t.id;
+    Trace.get t.id;
     x
   | Unresolved b ->
-    Suspend.enter (fun ctx enqueue ->
+    Suspend.enter "Promise.await" (fun ctx enqueue ->
         match Broadcast.suspend b (fun () -> enqueue (Ok ())) with
         | None -> ()  (* We got resumed immediately *)
         | Some request ->
@@ -53,15 +53,15 @@ let await t =
           | Unresolved _ ->
             (* We observed the promise to be still unresolved after registering a waiter.
                Therefore any resolution must happen after we were registered and we will be notified. *)
-            Trace.try_read t.id;
+            Trace.try_get t.id;
             Cancel.Fiber_context.set_cancel_fn ctx (fun ex ->
                 if Broadcast.cancel request then enqueue (Error ex)
                 (* else already resumed *)
               )
       );
     match Atomic.get t.state with
     | Resolved x ->
-      Trace.read t.id;
+      Trace.get t.id;
       x
     | Unresolved _ -> assert false
 
@@ -76,7 +76,7 @@ let try_resolve t v =
     | Resolved _ -> false
     | Unresolved b as prev ->
       if Atomic.compare_and_set t.state prev (Resolved v) then (
-        Trace.resolve t.id;
+        Trace.put t.id;
         Broadcast.resume_all b;
         true
       ) else (

lib_eio/core/single_waiter.ml

@@ -10,9 +10,9 @@ let create () = { wake = ignore }
 
 let wake t v = t.wake v
 
-let await t id =
+let await t op id =
   let x =
-    Suspend.enter @@ fun ctx enqueue ->
+    Suspend.enter op @@ fun ctx enqueue ->
     Cancel.Fiber_context.set_cancel_fn ctx (fun ex ->
         t.wake <- ignore;
         enqueue (Error ex)
@@ -23,5 +23,5 @@ let await t id =
         enqueue x
       )
   in
-  Trace.read id;
+  Trace.get id;
   x

lib_eio/core/suspend.ml

@@ -1,10 +1,12 @@
 type 'a enqueue = ('a, exn) result -> unit
 type _ Effect.t += Suspend : (Cancel.fiber_context -> 'a enqueue -> unit) -> 'a Effect.t
 
-let enter_unchecked fn = Effect.perform (Suspend fn)
+let enter_unchecked op fn =
+  Trace.suspend_fiber op;
+  Effect.perform (Suspend fn)
 
-let enter fn =
-  enter_unchecked @@ fun fiber enqueue ->
+let enter op fn =
+  enter_unchecked op @@ fun fiber enqueue ->
   match Cancel.Fiber_context.get_error fiber with
   | None -> fn fiber enqueue
   | Some ex -> enqueue (Error ex)