Merge pull request #19 from c-cube/wip-tracing-support

wip: tracing support

.github/workflows/main.yml

@@ -39,4 +39,6 @@ jobs:
         if: matrix.ocaml-compiler == '5.0'
       - run: opam exec -- dune build @install @runtest
         if: matrix.ocaml-compiler == '5.0'
+      - run: opam install trace thread-local-storage
+      - run: opam exec -- dune build @install
 

benchs/fib_rec.ml

@@ -1,4 +1,7 @@
 open Moonpool
+module Trace = Trace_core
+
+let ( let@ ) = ( @@ )
 
 let rec fib_direct x =
   if x <= 1 then
@@ -10,7 +13,7 @@ let cutoff = ref 20
 
 let rec fib ~on x : int Fut.t =
   if x <= !cutoff then
-    Fut.spawn ~on (fun () -> fib_direct x)
+    Fut.spawn ~name:"fib" ~on (fun () -> fib_direct x)
   else
     let open Fut.Infix in
     let+ t1 = fib ~on (x - 1) and+ t2 = fib ~on (x - 2) in
@@ -27,14 +30,14 @@ let fib_fj ~on x : int Fut.t =
       n1 + n2
     )
   in
-  Fut.spawn ~on (fun () -> fib_rec x)
+  Fut.spawn ~name:"fib" ~on (fun () -> fib_rec x)
 
 let fib_await ~on x : int Fut.t =
   let rec fib_rec x : int Fut.t =
     if x <= !cutoff then
-      Fut.spawn ~on (fun () -> fib_direct x)
+      Fut.spawn ~name:"fib" ~on (fun () -> fib_direct x)
     else
-      Fut.spawn ~on (fun () ->
+      Fut.spawn ~name:"fib" ~on (fun () ->
           let n1 = fib_rec (x - 1) in
           let n2 = fib_rec (x - 2) in
           let n1 = Fut.await n1 in
@@ -66,6 +69,7 @@ let str_of_int_opt = function
   | Some i -> Printf.sprintf "Some %d" i
 
 let run ~psize ~n ~seq ~dl ~fj ~await ~niter ~kind () : unit =
+  let@ _sp = Trace.with_span ~__FILE__ ~__LINE__ "fib.run" in
   let pool = lazy (create_pool ~kind ~psize ()) in
   let dl_pool =
     lazy
@@ -108,6 +112,7 @@ let run ~psize ~n ~seq ~dl ~fj ~await ~niter ~kind () : unit =
     Ws_pool.shutdown (Lazy.force pool)
 
 let () =
+  let@ () = Trace_tef.with_setup () in
   let n = ref 40 in
   let psize = ref None in
   let seq = ref false in

benchs/pi.ml

@@ -7,6 +7,7 @@ let j = ref 0
 let spf = Printf.sprintf
 
 let run_sequential (num_steps : int) : float =
+  let@ _sp = Trace.with_span ~__FILE__ ~__LINE__ "pi.seq" in
   let step = 1. /. float num_steps in
   let sum = ref 0. in
   for i = 0 to num_steps - 1 do
@@ -42,6 +43,11 @@ let run_par1 ~kind (num_steps : int) : float =
 
   (* one chunk of the work *)
   let run_task _idx_task : unit =
+    let@ _sp =
+      Trace.with_span ~__FILE__ ~__LINE__ "pi.slice" ~data:(fun () ->
+          [ "i", `Int _idx_task ])
+    in
+
     let sum = ref 0. in
     let i = ref 0 in
     while !i < num_steps do
@@ -69,7 +75,7 @@ let run_fork_join ~kind num_steps : float =
   let step = 1. /. float num_steps in
   let global_sum = Lock.create 0. in
 
-  Ws_pool.run_wait_block pool (fun () ->
+  Ws_pool.run_wait_block ~name:"pi.fj" pool (fun () ->
       Fork_join.for_
         ~chunk_size:(3 + (num_steps / num_tasks))
         num_steps
@@ -99,6 +105,9 @@ type mode =
 
 let () =
   let@ () = Trace_tef.with_setup () in
+  Trace.set_thread_name "main";
+  let@ _sp = Trace.with_span ~__FILE__ ~__LINE__ "main" in
+
   let mode = ref Sequential in
   let n = ref 1000 in
   let time = ref false in

dune-project

@@ -28,6 +28,7 @@
     (>= 1.9.0)
     :with-test)))
  (depopts
+   (trace (>= 0.6))
    thread-local-storage
    (domain-local-await (>= 0.2)))
  (tags

moonpool.opam

@@ -19,6 +19,7 @@ depends: [
   "mdx" {>= "1.9.0" & with-test}
 ]
 depopts: [
+  "trace" {>= "0.6"}
   "thread-local-storage"
   "domain-local-await" {>= "0.2"}
 ]

src/dune

@@ -1,7 +1,7 @@
 (library
  (public_name moonpool)
  (name moonpool)
- (private_modules d_pool_ dla_)
+ (private_modules d_pool_ dla_ tracing_)
  (preprocess
   (action
    (run %{project_root}/src/cpp/cpp.exe %{input-file})))
@@ -9,6 +9,9 @@
             (select thread_local_storage_.ml from
              (thread-local-storage -> thread_local_storage_.stub.ml)
              (-> thread_local_storage_.real.ml))
+            (select tracing_.ml from
+             (trace.core -> tracing_.real.ml)
+             (-> tracing_.dummy.ml))
             (select dla_.ml from
              (domain-local-await -> dla_.real.ml)
              ( -> dla_.dummy.ml))))

src/fifo_pool.ml

@@ -3,9 +3,14 @@ include Runner
 
 let ( let@ ) = ( @@ )
 
+type task_with_name = {
+  f: unit -> unit;
+  name: string;
+}
+
 type state = {
   threads: Thread.t array;
-  q: task Bb_queue.t;  (** Queue for tasks. *)
+  q: task_with_name Bb_queue.t;  (** Queue for tasks. *)
 }
 (** internal state *)
 
@@ -13,7 +18,7 @@ let[@inline] size_ (self : state) = Array.length self.threads
 let[@inline] num_tasks_ (self : state) : int = Bb_queue.size self.q
 
 (** Run [task] as is, on the pool. *)
-let schedule_ (self : state) (task : task) : unit =
+let schedule_ (self : state) (task : task_with_name) : unit =
   try Bb_queue.push self.q task with Bb_queue.Closed -> raise Shutdown
 
 type around_task = AT_pair : (t -> 'a) * (t -> 'a -> unit) -> around_task
@@ -22,19 +27,33 @@ let worker_thread_ (self : state) (runner : t) ~on_exn ~around_task : unit =
   TLS.get Runner.For_runner_implementors.k_cur_runner := Some runner;
   let (AT_pair (before_task, after_task)) = around_task in
 
-  let run_task task : unit =
+  let cur_span = ref Tracing_.dummy_span in
+
+  let[@inline] exit_span_ () =
+    Tracing_.exit_span !cur_span;
+    cur_span := Tracing_.dummy_span
+  in
+
+  let run_another_task ~name task' = schedule_ self { f = task'; name } in
+
+  let run_task (task : task_with_name) : unit =
     let _ctx = before_task runner in
+    cur_span := Tracing_.enter_span task.name;
     (* run the task now, catching errors *)
-    (try Suspend_.with_suspend task ~run:(fun task' -> schedule_ self task')
+    (try
+       Suspend_.with_suspend task.f ~name:task.name ~run:run_another_task
+         ~on_suspend:exit_span_
      with e ->
        let bt = Printexc.get_raw_backtrace () in
        on_exn e bt);
+    exit_span_ ();
     after_task runner _ctx
   in
 
   let main_loop () =
     let continue = ref true in
     while !continue do
+      assert (!cur_span = Tracing_.dummy_span);
       match Bb_queue.pop self.q with
       | task -> run_task task
       | exception Bb_queue.Closed -> continue := false
@@ -84,10 +103,12 @@ let create ?(on_init_thread = default_thread_init_exit_)
     { threads = Array.make num_threads dummy; q = Bb_queue.create () }
   in
 
+  let run_async ~name f = schedule_ pool { f; name } in
+
   let runner =
     Runner.For_runner_implementors.create
       ~shutdown:(fun ~wait () -> shutdown_ pool ~wait)
-      ~run_async:(fun f -> schedule_ pool f)
+      ~run_async
       ~size:(fun () -> size_ pool)
       ~num_tasks:(fun () -> num_tasks_ pool)
       ()

src/fork_join.ml

@@ -48,7 +48,7 @@ module State_ = struct
       Suspend_.suspend
         {
           Suspend_.handle =
-            (fun ~run:_ suspension ->
+            (fun ~name:_ ~run:_ suspension ->
               while
                 let old_st = A.get self in
                 match old_st with
@@ -113,19 +113,19 @@ let for_ ?chunk_size n (f : int -> int -> unit) : unit =
         max 1 (1 + (n / D_pool_.n_domains ()))
     in
 
-    let start_tasks ~run (suspension : Suspend_.suspension) =
+    let start_tasks ~name ~run (suspension : Suspend_.suspension) =
       let task_for ~offset ~len_range =
         match f offset (offset + len_range - 1) with
         | () ->
           if A.fetch_and_add missing (-len_range) = len_range then
             (* all tasks done successfully *)
-            suspension (Ok ())
+            run ~name (fun () -> suspension (Ok ()))
         | exception exn ->
           let bt = Printexc.get_raw_backtrace () in
           if not (A.exchange has_failed true) then
             (* first one to fail, and [missing] must be >= 2
                because we're not decreasing it. *)
-            suspension (Error (exn, bt))
+            run ~name (fun () -> suspension (Error (exn, bt)))
       in
 
       let i = ref 0 in
@@ -135,7 +135,7 @@ let for_ ?chunk_size n (f : int -> int -> unit) : unit =
         let len_range = min chunk_size (n - offset) in
         assert (offset + len_range <= n);
 
-        run (fun () -> task_for ~offset ~len_range);
+        run ~name (fun () -> task_for ~offset ~len_range);
         i := !i + len_range
       done
     in
@@ -143,9 +143,9 @@ let for_ ?chunk_size n (f : int -> int -> unit) : unit =
     Suspend_.suspend
       {
         Suspend_.handle =
-          (fun ~run suspension ->
+          (fun ~name ~run suspension ->
             (* run tasks, then we'll resume [suspension] *)
-            start_tasks ~run suspension);
+            start_tasks ~run ~name suspension);
       }
   )
 

src/fut.ml

@@ -5,16 +5,24 @@ type 'a waiter = 'a or_error -> unit
 
 type 'a state =
   | Done of 'a or_error
-  | Waiting of 'a waiter list
+  | Waiting of {
+      waiters: 'a waiter list;
+      name: string;
+    }
 
 type 'a t = { st: 'a state A.t } [@@unboxed]
 type 'a promise = 'a t
 
-let make () =
+let[@inline] get_name_ (self : _ t) =
-  let fut = { st = A.make (Waiting []) } in
+  match A.get self.st with
+  | Done _ -> ""
+  | Waiting { name; _ } -> name
+
+let make ?(name = "") () =
+  let fut = { st = A.make (Waiting { waiters = []; name }) } in
   fut, fut
 
-let of_result x : _ t = { st = A.make (Done x) }
+let[@inline] of_result x : _ t = { st = A.make (Done x) }
 let[@inline] return x : _ t = of_result (Ok x)
 let[@inline] fail e bt : _ t = of_result (Error (e, bt))
 
@@ -53,9 +61,8 @@ let on_result (self : _ t) (f : _ waiter) : unit =
     | Done x ->
       f x;
       false
-    | Waiting l ->
+    | Waiting { waiters = l; name } ->
-      let must_retry = not (A.compare_and_set self.st st (Waiting (f :: l))) in
+      not (A.compare_and_set self.st st (Waiting { waiters = f :: l; name }))
-      must_retry
   do
     Domain_.relax ()
   done
@@ -63,28 +70,31 @@ let on_result (self : _ t) (f : _ waiter) : unit =
 exception Already_fulfilled
 
 let fulfill (self : _ t) (r : _ result) : unit =
+  let fs = ref [] in
   while
     let st = A.get self.st in
     match st with
     | Done _ -> raise Already_fulfilled
-    | Waiting l ->
+    | Waiting { waiters = l; name = _ } ->
       let did_swap = A.compare_and_set self.st st (Done r) in
       if did_swap then (
         (* success, now call all the waiters *)
-        List.iter (fun f -> try f r with _ -> ()) l;
+        fs := l;
         false
       ) else
         true
   do
     Domain_.relax ()
-  done
+  done;
+  List.iter (fun f -> try f r with _ -> ()) !fs;
+  ()
 
 let[@inline] fulfill_idempotent self r =
   try fulfill self r with Already_fulfilled -> ()
 
 (* ### combinators ### *)
 
-let spawn ~on f : _ t =
+let spawn ?name ~on f : _ t =
   let fut, promise = make () in
 
   let task () =
@@ -97,13 +107,13 @@ let spawn ~on f : _ t =
     fulfill promise res
   in
 
-  Runner.run_async on task;
+  Runner.run_async ?name on task;
   fut
 
-let spawn_on_current_runner f : _ t =
+let spawn_on_current_runner ?name f : _ t =
   match Runner.get_current_runner () with
   | None -> failwith "Fut.spawn_on_current_runner: not running on a runner"
-  | Some on -> spawn ~on f
+  | Some on -> spawn ?name ~on f
 
 let reify_error (f : 'a t) : 'a or_error t =
   match peek f with
@@ -113,7 +123,7 @@ let reify_error (f : 'a t) : 'a or_error t =
     on_result f (fun r -> fulfill promise (Ok r));
     fut
 
-let get_runner_ ?on () : Runner.t option =
+let[@inline] get_runner_ ?on () : Runner.t option =
   match on with
   | Some _ as r -> r
   | None -> Runner.get_current_runner ()
@@ -129,20 +139,22 @@ let map ?on ~f fut : _ t =
     | Error e_bt -> Error e_bt
   in
 
+  let name = get_name_ fut in
   match peek fut, get_runner_ ?on () with
   | Some res, None -> of_result @@ map_immediate_ res
   | Some res, Some runner ->
-    let fut2, promise = make () in
+    let fut2, promise = make ~name () in
-    Runner.run_async runner (fun () -> fulfill promise @@ map_immediate_ res);
+    Runner.run_async ~name runner (fun () ->
+        fulfill promise @@ map_immediate_ res);
     fut2
   | None, None ->
-    let fut2, promise = make () in
+    let fut2, promise = make ~name () in
     on_result fut (fun res -> fulfill promise @@ map_immediate_ res);
     fut2
   | None, Some runner ->
-    let fut2, promise = make () in
+    let fut2, promise = make ~name () in
     on_result fut (fun res ->
-        Runner.run_async runner (fun () ->
+        Runner.run_async ~name runner (fun () ->
             fulfill promise @@ map_immediate_ res));
     fut2
 
@@ -151,7 +163,7 @@ let join (fut : 'a t t) : 'a t =
   | Some (Ok f) -> f
   | Some (Error (e, bt)) -> fail e bt
   | None ->
-    let fut2, promise = make () in
+    let fut2, promise = make ~name:(get_name_ fut) () in
     on_result fut (function
       | Ok sub_fut -> on_result sub_fut (fulfill promise)
       | Error _ as e -> fulfill promise e);
@@ -174,19 +186,20 @@ let bind ?on ~f fut : _ t =
     on_result f_res_fut (fun r -> fulfill promise r)
   in
 
+  let name = get_name_ fut in
   match peek fut, get_runner_ ?on () with
   | Some res, Some runner ->
-    let fut2, promise = make () in
+    let fut2, promise = make ~name () in
-    Runner.run_async runner (bind_and_fulfill res promise);
+    Runner.run_async ~name runner (bind_and_fulfill res promise);
     fut2
   | Some res, None -> apply_f_to_res res
   | None, Some runner ->
-    let fut2, promise = make () in
+    let fut2, promise = make ~name () in
     on_result fut (fun r ->
-        Runner.run_async runner (bind_and_fulfill r promise));
+        Runner.run_async ~name runner (bind_and_fulfill r promise));
     fut2
   | None, None ->
-    let fut2, promise = make () in
+    let fut2, promise = make ~name () in
     on_result fut (fun res -> bind_and_fulfill res promise ());
     fut2
 
@@ -210,7 +223,7 @@ let both a b : _ t =
   | Some (Ok x), Some (Ok y) -> return (x, y)
   | Some (Error (e, bt)), _ | _, Some (Error (e, bt)) -> fail e bt
   | _ ->
-    let fut, promise = make () in
+    let fut, promise = make ~name:(get_name_ a) () in
 
     let st = A.make `Neither in
     on_result a (function
@@ -243,7 +256,7 @@ let choose a b : _ t =
   | _, Some (Ok y) -> return (Either.Right y)
   | Some (Error (e, bt)), Some (Error _) -> fail e bt
   | _ ->
-    let fut, promise = make () in
+    let fut, promise = make ~name:(get_name_ a) () in
 
     let one_failure = A.make false in
     on_result a (function
@@ -266,7 +279,7 @@ let choose_same a b : _ t =
   | _, Some (Ok y) -> return y
   | Some (Error (e, bt)), Some (Error _) -> fail e bt
   | _ ->
-    let fut, promise = make () in
+    let fut, promise = make ~name:(get_name_ a) () in
 
     let one_failure = A.make false in
     on_result a (function
@@ -413,9 +426,11 @@ let await (fut : 'a t) : 'a =
     Suspend_.suspend
       {
         Suspend_.handle =
-          (fun ~run k ->
+          (fun ~name ~run k ->
             on_result fut (function
-              | Ok _ -> run (fun () -> k (Ok ()))
+              | Ok _ ->
+                (* schedule continuation with the same name *)
+                run ~name (fun () -> k (Ok ()))
               | Error (exn, bt) ->
                 (* fail continuation immediately *)
                 k (Error (exn, bt))));

src/fut.mli

@@ -26,8 +26,9 @@ type 'a promise
 (** A promise, which can be fulfilled exactly once to set
       the corresponding future *)
 
-val make : unit -> 'a t * 'a promise
+val make : ?name:string -> unit -> 'a t * 'a promise
-(** Make a new future with the associated promise *)
+(** Make a new future with the associated promise.
+    @param name name for the future, used for tracing. since NEXT_RELEASE. *)
 
 val on_result : 'a t -> ('a or_error -> unit) -> unit
 (** [on_result fut f] registers [f] to be called in the future
@@ -81,11 +82,11 @@ val is_done : _ t -> bool
 
 (** {2 Combinators} *)
 
-val spawn : on:Runner.t -> (unit -> 'a) -> 'a t
+val spawn : ?name:string -> on:Runner.t -> (unit -> 'a) -> 'a t
 (** [spaw ~on f] runs [f()] on the given runner [on], and return a future that will
       hold its result. *)
 
-val spawn_on_current_runner : (unit -> 'a) -> 'a t
+val spawn_on_current_runner : ?name:string -> (unit -> 'a) -> 'a t
 (** This must be run from inside a runner, and schedules
     the new task on it as well.
 

src/immediate_runner.ml

@@ -1,9 +1,19 @@
 include Runner
 
+let run_async_ ~name f =
+  let sp = Tracing_.enter_span name in
+  try
+    let x = f () in
+    Tracing_.exit_span sp;
+    x
+  with e ->
+    let bt = Printexc.get_raw_backtrace () in
+    Tracing_.exit_span sp;
+    Printexc.raise_with_backtrace e bt
+
 let runner : t =
   Runner.For_runner_implementors.create
     ~size:(fun () -> 0)
     ~num_tasks:(fun () -> 0)
     ~shutdown:(fun ~wait:_ () -> ())
-    ~run_async:(fun f -> f ())
+    ~run_async:run_async_ ()
-    ()

src/moonpool.mli

@@ -23,10 +23,13 @@ val start_thread_on_some_domain : ('a -> unit) -> 'a -> Thread.t
     to run the thread. This ensures that we don't always pick the same domain
     to run all the various threads needed in an application (timers, event loops, etc.) *)
 
-val run_async : Runner.t -> (unit -> unit) -> unit
+val run_async : ?name:string -> Runner.t -> (unit -> unit) -> unit
 (** [run_async runner task] schedules the task to run
   on the given runner. This means [task()] will be executed
   at some point in the future, possibly in another thread.
+  @param name if provided and [Trace] is present in dependencies, a span
+    will be created when the task starts, and will stop when the task is over.
+    (since NEXT_RELEASE)
   @since 0.5 *)
 
 val recommended_thread_count : unit -> int
@@ -35,13 +38,16 @@ val recommended_thread_count : unit -> int
   this because many of them will be blocked most of the time).
   @since 0.5 *)
 
-val spawn : on:Runner.t -> (unit -> 'a) -> 'a Fut.t
+val spawn : ?name:string -> on:Runner.t -> (unit -> 'a) -> 'a Fut.t
 (** [spawn ~on f] runs [f()] on the runner (a thread pool typically)
     and returns a future result for it. See {!Fut.spawn}.
+    @param name if provided and [Trace] is present in dependencies,
+      a span will be created for the future. (since NEXT_RELEASE)
     @since 0.5 *)
 
-val spawn_on_current_runner : (unit -> 'a) -> 'a Fut.t
+val spawn_on_current_runner : ?name:string -> (unit -> 'a) -> 'a Fut.t
 (** See {!Fut.spawn_on_current_runner}.
+    @param name see {!spawn}. since NEXT_RELEASE.
     @since 0.5 *)
 
 [@@@ifge 5.0]

src/runner.ml

@@ -3,7 +3,7 @@ module TLS = Thread_local_storage_
 type task = unit -> unit
 
 type t = {
-  run_async: (unit -> unit) -> unit;
+  run_async: name:string -> task -> unit;
   shutdown: wait:bool -> unit -> unit;
   size: unit -> int;
   num_tasks: unit -> int;
@@ -11,7 +11,7 @@ type t = {
 
 exception Shutdown
 
-let[@inline] run_async (self : t) f : unit = self.run_async f
+let[@inline] run_async ?(name = "") (self : t) f : unit = self.run_async ~name f
 let[@inline] shutdown (self : t) : unit = self.shutdown ~wait:true ()
 
 let[@inline] shutdown_without_waiting (self : t) : unit =
@@ -20,9 +20,9 @@ let[@inline] shutdown_without_waiting (self : t) : unit =
 let[@inline] num_tasks (self : t) : int = self.num_tasks ()
 let[@inline] size (self : t) : int = self.size ()
 
-let run_wait_block self (f : unit -> 'a) : 'a =
+let run_wait_block ?name self (f : unit -> 'a) : 'a =
   let q = Bb_queue.create () in
-  run_async self (fun () ->
+  run_async ?name self (fun () ->
       try
         let x = f () in
         Bb_queue.push q (Ok x)

src/runner.mli

@@ -33,13 +33,16 @@ val shutdown_without_waiting : t -> unit
 
 exception Shutdown
 
-val run_async : t -> task -> unit
+val run_async : ?name:string -> t -> task -> unit
 (** [run_async pool f] schedules [f] for later execution on the runner
     in one of the threads. [f()] will run on one of the runner's
     worker threads/domains.
+    @param name if provided and [Trace] is present in dependencies, a span
+    will be created when the task starts, and will stop when the task is over.
+    (since NEXT_RELEASE)
     @raise Shutdown if the runner was shut down before [run_async] was called. *)
 
-val run_wait_block : t -> (unit -> 'a) -> 'a
+val run_wait_block : ?name:string -> t -> (unit -> 'a) -> 'a
 (** [run_wait_block pool f] schedules [f] for later execution
     on the pool, like {!run_async}.
     It then blocks the current thread until [f()] is done executing,
@@ -47,7 +50,10 @@ val run_wait_block : t -> (unit -> 'a) -> 'a
     will raise it as well.
 
     {b NOTE} be careful with deadlocks (see notes in {!Fut.wait_block}
-      about the required discipline to avoid deadlocks). *)
+      about the required discipline to avoid deadlocks).
+    @raise Shutdown if the runner was already shut down *)
+
+(** {2 Implementing runners} *)
 
 (** This module is specifically intended for users who implement their
     own runners. Regular users of Moonpool should not need to look at it. *)
@@ -56,7 +62,7 @@ module For_runner_implementors : sig
     size:(unit -> int) ->
     num_tasks:(unit -> int) ->
     shutdown:(wait:bool -> unit -> unit) ->
-    run_async:(task -> unit) ->
+    run_async:(name:string -> task -> unit) ->
     unit ->
     t
   (** Create a new runner.
@@ -65,6 +71,10 @@ module For_runner_implementors : sig
       so that {!Fork_join} and other 5.x features work properly. *)
 
   val k_cur_runner : t option ref Thread_local_storage_.key
+  (** Key that should be used by each runner to store itself in TLS
+      on every thread it controls, so that tasks running on these threads
+      can access the runner. This is necessary for {!get_current_runner}
+      to work. *)
 end
 
 val get_current_runner : unit -> t option

src/suspend_.ml

@@ -1,7 +1,9 @@
 type suspension = (unit, exn * Printexc.raw_backtrace) result -> unit
 type task = unit -> unit
 
-type suspension_handler = { handle: run:(task -> unit) -> suspension -> unit }
+type suspension_handler = {
+  handle: name:string -> run:(name:string -> task -> unit) -> suspension -> unit;
+}
 [@@unboxed]
 
 [@@@ifge 5.0]
@@ -13,7 +15,8 @@ type _ Effect.t += Suspend : suspension_handler -> unit Effect.t
 
 let[@inline] suspend h = Effect.perform (Suspend h)
 
-let with_suspend ~(run : task -> unit) (f : unit -> unit) : unit =
+let with_suspend ~name ~on_suspend ~(run : name:string -> task -> unit)
+    (f : unit -> unit) : unit =
   let module E = Effect.Deep in
   (* effect handler *)
   let effc : type e. e Effect.t -> ((e, _) E.continuation -> _) option =
@@ -21,11 +24,12 @@ let with_suspend ~(run : task -> unit) (f : unit -> unit) : unit =
     | Suspend h ->
       Some
         (fun k ->
+          on_suspend ();
           let k' : suspension = function
             | Ok () -> E.continue k ()
             | Error (exn, bt) -> E.discontinue_with_backtrace k exn bt
           in
-          h.handle ~run k')
+          h.handle ~name ~run k')
     | _ -> None
   in
 
@@ -34,14 +38,16 @@ let with_suspend ~(run : task -> unit) (f : unit -> unit) : unit =
 (* DLA interop *)
 let prepare_for_await () : Dla_.t =
   (* current state *)
-  let st : ((task -> unit) * suspension) option A.t = A.make None in
+  let st : (string * (name:string -> task -> unit) * suspension) option A.t =
+    A.make None
+  in
 
   let release () : unit =
     match A.exchange st None with
     | None -> ()
-    | Some (run, k) -> run (fun () -> k (Ok ()))
+    | Some (name, run, k) -> run ~name (fun () -> k (Ok ()))
   and await () : unit =
-    suspend { handle = (fun ~run k -> A.set st (Some (run, k))) }
+    suspend { handle = (fun ~name ~run k -> A.set st (Some (name, run, k))) }
   in
 
   let t = { Dla_.release; await } in
@@ -50,7 +56,7 @@ let prepare_for_await () : Dla_.t =
 [@@@ocaml.alert "+unstable"]
 [@@@else_]
 
-let[@inline] with_suspend ~run:_ f = f ()
+let[@inline] with_suspend ~name:_ ~on_suspend:_ ~run:_ f = f ()
 let[@inline] prepare_for_await () = { Dla_.release = ignore; await = ignore }
 
 [@@@endif]

src/suspend_.mli

@@ -8,12 +8,15 @@ type suspension = (unit, exn * Printexc.raw_backtrace) result -> unit
 
 type task = unit -> unit
 
-type suspension_handler = { handle: run:(task -> unit) -> suspension -> unit }
+type suspension_handler = {
+  handle: name:string -> run:(name:string -> task -> unit) -> suspension -> unit;
+}
 [@@unboxed]
 (** The handler that knows what to do with the suspended computation.
 
-   The handler is given two things:
+   The handler is given a few things:
 
+   - the name (if any) of the current computation
    - the suspended computation (which can be resumed with a result
      eventually);
    - a [run] function that can be used to start tasks to perform some
@@ -51,10 +54,18 @@ val suspend : suspension_handler -> unit
 val prepare_for_await : unit -> Dla_.t
 (** Our stub for DLA. Unstable. *)
 
-val with_suspend : run:(task -> unit) -> (unit -> unit) -> unit
+val with_suspend :
+  name:string ->
+  on_suspend:(unit -> unit) ->
+  run:(name:string -> task -> unit) ->
+  (unit -> unit) ->
+  unit
 (** [with_suspend ~run f] runs [f()] in an environment where [suspend]
     will work. If [f()] suspends with suspension handler [h],
     this calls [h ~run k] where [k] is the suspension.
+    The suspension should always run in a new task, via [run].
+
+    @param on_suspend called when [f()] suspends itself.
 
     This will not do anything on OCaml 4.x.
 *)

src/tracing_.dummy.ml

@@ -0,0 +1,4 @@
+let enabled () = false
+let dummy_span = 0L
+let enter_span _name = dummy_span
+let exit_span = ignore

src/tracing_.mli

@@ -0,0 +1,4 @@
+val dummy_span : int64
+val enter_span : string -> int64
+val exit_span : int64 -> unit
+val enabled : unit -> bool

src/tracing_.real.ml

@@ -0,0 +1,13 @@
+module Trace = Trace_core
+
+let enabled = Trace.enabled
+let dummy_span = Int64.min_int
+let dummy_file_ = "<unknown file>"
+
+let[@inline] enter_span name : int64 =
+  if name = "" then
+    dummy_span
+  else
+    Trace.enter_span ~__FILE__:dummy_file_ ~__LINE__:0 name
+
+let[@inline] exit_span sp = if sp <> dummy_span then Trace.exit_span sp

src/ws_pool.ml

@@ -13,10 +13,16 @@ module Id = struct
   let equal : t -> t -> bool = ( == )
 end
 
+type task_with_name = {
+  f: task;
+  name: string;
+}
+
 type worker_state = {
   pool_id_: Id.t;  (** Unique per pool *)
   mutable thread: Thread.t;
-  q: task WSQ.t;  (** Work stealing queue *)
+  q: task_with_name WSQ.t;  (** Work stealing queue *)
+  mutable cur_span: int64;
   rng: Random.State.t;
 }
 (** State for a given worker. Only this worker is
@@ -29,7 +35,8 @@ type state = {
   id_: Id.t;
   active: bool A.t;  (** Becomes [false] when the pool is shutdown. *)
   workers: worker_state array;  (** Fixed set of workers. *)
-  main_q: task Queue.t;  (** Main queue for tasks coming from the outside *)
+  main_q: task_with_name Queue.t;
+      (** Main queue for tasks coming from the outside *)
   mutable n_waiting: int; (* protected by mutex *)
   mutable n_waiting_nonzero: bool;  (** [n_waiting > 0] *)
   mutex: Mutex.t;
@@ -65,9 +72,10 @@ let[@inline] try_wake_someone_ (self : state) : unit =
   )
 
 (** Run [task] as is, on the pool. *)
-let schedule_task_ (self : state) (w : worker_state option) (task : task) : unit
+let schedule_task_ (self : state) ~name (w : worker_state option) (f : task) :
-    =
+    unit =
   (* Printf.printf "schedule task now (%d)\n%!" (Thread.id @@ Thread.self ()); *)
+  let task = { f; name } in
   match w with
   | Some w when Id.equal self.id_ w.pool_id_ ->
     (* we're on this same pool, schedule in the worker's state. Otherwise
@@ -96,24 +104,36 @@ let schedule_task_ (self : state) (w : worker_state option) (task : task) : unit
       raise Shutdown
 
 (** Run this task, now. Must be called from a worker. *)
-let run_task_now_ (self : state) ~runner task : unit =
+let run_task_now_ (self : state) ~runner (w : worker_state) ~name task : unit =
   (* Printf.printf "run task now (%d)\n%!" (Thread.id @@ Thread.self ()); *)
   let (AT_pair (before_task, after_task)) = self.around_task in
   let _ctx = before_task runner in
+
+  w.cur_span <- Tracing_.enter_span name;
+  let[@inline] exit_span_ () =
+    Tracing_.exit_span w.cur_span;
+    w.cur_span <- Tracing_.dummy_span
+  in
+
+  let run_another_task ~name task' =
+    let w = find_current_worker_ () in
+    schedule_task_ self w ~name task'
+  in
+
   (* run the task now, catching errors *)
   (try
      (* run [task()] and handle [suspend] in it *)
-     Suspend_.with_suspend task ~run:(fun task' ->
+     Suspend_.with_suspend task ~name ~run:run_another_task
-         let w = find_current_worker_ () in
+       ~on_suspend:exit_span_
-         schedule_task_ self w task')
    with e ->
      let bt = Printexc.get_raw_backtrace () in
      self.on_exn e bt);
+  exit_span_ ();
   after_task runner _ctx
 
-let[@inline] run_async_ (self : state) (task : task) : unit =
+let[@inline] run_async_ (self : state) ~name (f : task) : unit =
   let w = find_current_worker_ () in
-  schedule_task_ self w task
+  schedule_task_ self w ~name f
 
 (* TODO: function to schedule many tasks from the outside.
     - build a queue
@@ -122,8 +142,6 @@ let[@inline] run_async_ (self : state) (task : task) : unit =
     - wakeup all (broadcast)
     - unlock *)
 
-let run = run_async
-
 (** Wait on condition. Precondition: we hold the mutex. *)
 let[@inline] wait_ (self : state) : unit =
   self.n_waiting <- self.n_waiting + 1;
@@ -132,10 +150,11 @@ let[@inline] wait_ (self : state) : unit =
   self.n_waiting <- self.n_waiting - 1;
   if self.n_waiting = 0 then self.n_waiting_nonzero <- false
 
-exception Got_task of task
+exception Got_task of task_with_name
 
 (** Try to steal a task *)
-let try_to_steal_work_once_ (self : state) (w : worker_state) : task option =
+let try_to_steal_work_once_ (self : state) (w : worker_state) :
+    task_with_name option =
   let init = Random.State.int w.rng (Array.length self.workers) in
 
   try
@@ -160,7 +179,7 @@ let worker_run_self_tasks_ (self : state) ~runner w : unit =
     match WSQ.pop w.q with
     | Some task ->
       try_wake_someone_ self;
-      run_task_now_ self ~runner task
+      run_task_now_ self ~runner w ~name:task.name task.f
     | None -> continue := false
   done
 
@@ -173,7 +192,7 @@ let worker_thread_ (self : state) ~(runner : t) (w : worker_state) : unit =
     worker_run_self_tasks_ self ~runner w;
     try_steal ()
   and run_task task : unit =
-    run_task_now_ self ~runner task;
+    run_task_now_ self ~runner w ~name:task.name task.f;
     main ()
   and try_steal () =
     match try_to_steal_work_once_ self w with
@@ -231,7 +250,7 @@ type ('a, 'b) create_args =
   'a
 (** Arguments used in {!create}. See {!create} for explanations. *)
 
-let dummy_task_ () = assert false
+let dummy_task_ = { f = ignore; name = "DUMMY_TASK" }
 
 let create ?(on_init_thread = default_thread_init_exit_)
     ?(on_exit_thread = default_thread_init_exit_) ?(on_exn = fun _ _ -> ())
@@ -256,6 +275,7 @@ let create ?(on_init_thread = default_thread_init_exit_)
         {
           pool_id_;
           thread = dummy;
+          cur_span = Tracing_.dummy_span;
           q = WSQ.create ~dummy:dummy_task_ ();
           rng = Random.State.make [| i |];
         })
@@ -279,7 +299,7 @@ let create ?(on_init_thread = default_thread_init_exit_)
   let runner =
     Runner.For_runner_implementors.create
       ~shutdown:(fun ~wait () -> shutdown_ pool ~wait)
-      ~run_async:(fun f -> run_async_ pool f)
+      ~run_async:(fun ~name f -> run_async_ pool ~name f)
       ~size:(fun () -> size_ pool)
       ~num_tasks:(fun () -> num_tasks_ pool)
       ()

src/ws_pool.mli

@@ -53,7 +53,3 @@ val with_ : (unit -> (t -> 'a) -> 'a, _) create_args
 
     Most parameters are the same as in {!create}.
     @since 0.3 *)
-
-val run : t -> (unit -> unit) -> unit
-  [@@deprecated "use run_async"]
-(** deprecated alias to {!run_async} *)

test/effect-based/t_fork_join_heavy.ml

@@ -33,10 +33,12 @@ let run ~min () =
       let l1, l2 =
         Fork_join.both
           (fun () ->
+            let@ _sp = Trace.with_span ~__FILE__ ~__LINE__ "fj.left" in
             Fork_join.map_list ~chunk_size
               (Fork_join.map_list ~chunk_size neg)
               l)
           (fun () ->
+            let@ _sp = Trace.with_span ~__FILE__ ~__LINE__ "fj.right" in
             Fork_join.map_list ~chunk_size
               (Fork_join.map_list ~chunk_size neg)
               ref_l1)