Makefile

@@ -67,14 +67,6 @@ bench-pi:
 		'./_build/default/benchs/pi.exe -n $(PI_NSTEPS) -j 16 -mode forkjoin -kind=pool' \
 		'./_build/default/benchs/pi.exe -n $(PI_NSTEPS) -j 20 -mode forkjoin -kind=pool'
 
-bench-repro-41:
-	dune build $(DUNE_OPTS_BENCH) examples/repro_41/run.exe
-	hyperfine --warmup=1 \
-		"./_build/default/examples/repro_41/run.exe 4 domainslib" \
-		"./_build/default/examples/repro_41/run.exe 4 moonpool" \
-		"./_build/default/examples/repro_41/run.exe 5 moonpool" \
-		"./_build/default/examples/repro_41/run.exe 5 seq"
-
 .PHONY: test clean bench-fib bench-pi
 
 VERSION=$(shell awk '/^version:/ {print $$2}' moonpool.opam)

examples/repro_41/dune

@@ -1,5 +0,0 @@
-(executables
- (names run)
- (enabled_if
-  (>= %{ocaml_version} 5.0))
- (libraries moonpool trace trace-tef domainslib))

examples/repro_41/run.ml

@@ -1,54 +0,0 @@
-(* fibo.ml *)
-let cutoff = 25
-let input = 40
-
-let rec fibo_seq n =
-  if n <= 1 then
-    n
-  else
-    fibo_seq (n - 1) + fibo_seq (n - 2)
-
-let rec fibo_domainslib ctx n =
-  if n <= cutoff then
-    fibo_seq n
-  else
-    let open Domainslib in
-    let fut1 = Task.async ctx (fun () -> fibo_domainslib ctx (n - 1)) in
-    let fut2 = Task.async ctx (fun () -> fibo_domainslib ctx (n - 2)) in
-    Task.await ctx fut1 + Task.await ctx fut2
-
-let rec fibo_moonpool ctx n =
-  if n <= cutoff then
-    fibo_seq n
-  else
-    let open Moonpool in
-    let fut1 = Fut.spawn ~on:ctx (fun () -> fibo_moonpool ctx (n - 1)) in
-    let fut2 = Fut.spawn ~on:ctx (fun () -> fibo_moonpool ctx (n - 2)) in
-    Fut.await fut1 + Fut.await fut2
-
-let usage =
-  "fibo.exe <num_domains> [ domainslib | moonpool | moonpool_fifo | seq ]"
-
-let num_domains = try int_of_string Sys.argv.(1) with _ -> failwith usage
-let implem = try Sys.argv.(2) with _ -> failwith usage
-
-let () =
-  let output =
-    match implem with
-    | "moonpool" ->
-      let open Moonpool in
-      let ctx = Ws_pool.create ~num_threads:num_domains () in
-      Ws_pool.run_wait_block ctx (fun () -> fibo_moonpool ctx input)
-    | "moonpool_fifo" ->
-      let open Moonpool in
-      let ctx = Fifo_pool.create ~num_threads:num_domains () in
-      Ws_pool.run_wait_block ctx (fun () -> fibo_moonpool ctx input)
-    | "domainslib" ->
-      let open Domainslib in
-      let pool = Task.setup_pool ~num_domains () in
-      Task.run pool (fun () -> fibo_domainslib pool input)
-    | "seq" -> fibo_seq input
-    | _ -> failwith usage
-  in
-  print_int output;
-  print_newline ()

src/core/background_thread.ml

@@ -6,15 +6,18 @@ type ('a, 'b) create_args =
   ?on_init_thread:(dom_id:int -> t_id:int -> unit -> unit) ->
   ?on_exit_thread:(dom_id:int -> t_id:int -> unit -> unit) ->
   ?on_exn:(exn -> Printexc.raw_backtrace -> unit) ->
+  ?around_task:(t -> 'b) * (t -> 'b -> unit) ->
   ?name:string ->
   'a
 (** Arguments used in {!create}. See {!create} for explanations. *)
 
-let create ?on_init_thread ?on_exit_thread ?on_exn ?name () : t =
-  Fifo_pool.create ?on_init_thread ?on_exit_thread ?on_exn ?name ~num_threads:1
-    ()
+let create ?on_init_thread ?on_exit_thread ?on_exn ?around_task ?name () : t =
+  Fifo_pool.create ?on_init_thread ?on_exit_thread ?on_exn ?around_task ?name
+    ~num_threads:1 ()
 
-let with_ ?on_init_thread ?on_exit_thread ?on_exn ?name () f =
-  let pool = create ?on_init_thread ?on_exit_thread ?on_exn ?name () in
+let with_ ?on_init_thread ?on_exit_thread ?on_exn ?around_task ?name () f =
+  let pool =
+    create ?on_init_thread ?on_exit_thread ?on_exn ?around_task ?name ()
+  in
   let@ () = Fun.protect ~finally:(fun () -> shutdown pool) in
   f pool

src/core/background_thread.mli

@@ -13,6 +13,7 @@ type ('a, 'b) create_args =
   ?on_init_thread:(dom_id:int -> t_id:int -> unit -> unit) ->
   ?on_exit_thread:(dom_id:int -> t_id:int -> unit -> unit) ->
   ?on_exn:(exn -> Printexc.raw_backtrace -> unit) ->
+  ?around_task:(t -> 'b) * (t -> 'b -> unit) ->
   ?name:string ->
   'a
 (** Arguments used in {!create}. See {!create} for explanations. *)

src/core/fifo_pool.ml

@@ -10,6 +10,7 @@ let ( let@ ) = ( @@ )
 type state = {
   threads: Thread.t array;
   q: task_full Bb_queue.t;  (** Queue for tasks. *)
+  around_task: WL.around_task;
   mutable as_runner: t;
   (* init options *)
   name: string option;
@@ -27,6 +28,11 @@ type worker_state = {
 
 let[@inline] size_ (self : state) = Array.length self.threads
 let[@inline] num_tasks_ (self : state) : int = Bb_queue.size self.q
+
+(*
+get_thread_state = TLS.get_opt k_worker_state
+  *)
+
 let default_thread_init_exit_ ~dom_id:_ ~t_id:_ () = ()
 
 let shutdown_ ~wait (self : state) : unit =
@@ -37,10 +43,13 @@ type ('a, 'b) create_args =
   ?on_init_thread:(dom_id:int -> t_id:int -> unit -> unit) ->
   ?on_exit_thread:(dom_id:int -> t_id:int -> unit -> unit) ->
   ?on_exn:(exn -> Printexc.raw_backtrace -> unit) ->
+  ?around_task:(t -> 'b) * (t -> 'b -> unit) ->
   ?num_threads:int ->
   ?name:string ->
   'a
 
+let default_around_task_ : WL.around_task = AT_pair (ignore, fun _ _ -> ())
+
 (** Run [task] as is, on the pool. *)
 let schedule_ (self : state) (task : task_full) : unit =
   try Bb_queue.push self.q task with Bb_queue.Closed -> raise Shutdown
@@ -79,6 +88,7 @@ let cleanup (self : worker_state) : unit =
 
 let worker_ops : worker_state WL.ops =
   let runner (st : worker_state) = st.st.as_runner in
+  let around_task st = st.st.around_task in
   let on_exn (st : worker_state) (ebt : Exn_bt.t) =
     st.st.on_exn (Exn_bt.exn ebt) (Exn_bt.bt ebt)
   in
@@ -86,6 +96,7 @@ let worker_ops : worker_state WL.ops =
     WL.schedule = schedule_w;
     runner;
     get_next_task;
+    around_task;
     on_exn;
     before_start;
     cleanup;
@@ -93,11 +104,19 @@ let worker_ops : worker_state WL.ops =
 
 let create_ ?(on_init_thread = default_thread_init_exit_)
     ?(on_exit_thread = default_thread_init_exit_) ?(on_exn = fun _ _ -> ())
-    ~threads ?name () : state =
+    ?around_task ~threads ?name () : state =
+  (* wrapper *)
+  let around_task =
+    match around_task with
+    | Some (f, g) -> WL.AT_pair (f, g)
+    | None -> default_around_task_
+  in
+
   let self =
     {
       threads;
       q = Bb_queue.create ();
+      around_task;
       as_runner = Runner.dummy;
       name;
       on_init_thread;
@@ -108,7 +127,8 @@ let create_ ?(on_init_thread = default_thread_init_exit_)
   self.as_runner <- runner_of_state self;
   self
 
-let create ?on_init_thread ?on_exit_thread ?on_exn ?num_threads ?name () : t =
+let create ?on_init_thread ?on_exit_thread ?on_exn ?around_task ?num_threads
+    ?name () : t =
   let num_domains = Domain_pool_.max_number_of_domains () in
 
   (* number of threads to run *)
@@ -120,7 +140,8 @@ let create ?on_init_thread ?on_exit_thread ?on_exn ?num_threads ?name () : t =
   let pool =
     let dummy_thread = Thread.self () in
     let threads = Array.make num_threads dummy_thread in
-    create_ ?on_init_thread ?on_exit_thread ?on_exn ~threads ?name ()
+    create_ ?on_init_thread ?on_exit_thread ?on_exn ?around_task ~threads ?name
+      ()
   in
   let runner = runner_of_state pool in
 
@@ -160,9 +181,11 @@ let create ?on_init_thread ?on_exit_thread ?on_exn ?num_threads ?name () : t =
 
   runner
 
-let with_ ?on_init_thread ?on_exit_thread ?on_exn ?num_threads ?name () f =
+let with_ ?on_init_thread ?on_exit_thread ?on_exn ?around_task ?num_threads
+    ?name () f =
   let pool =
-    create ?on_init_thread ?on_exit_thread ?on_exn ?num_threads ?name ()
+    create ?on_init_thread ?on_exit_thread ?on_exn ?around_task ?num_threads
+      ?name ()
   in
   let@ () = Fun.protect ~finally:(fun () -> shutdown pool) in
   f pool

src/core/fifo_pool.mli

@@ -20,6 +20,7 @@ type ('a, 'b) create_args =
   ?on_init_thread:(dom_id:int -> t_id:int -> unit -> unit) ->
   ?on_exit_thread:(dom_id:int -> t_id:int -> unit -> unit) ->
   ?on_exn:(exn -> Printexc.raw_backtrace -> unit) ->
+  ?around_task:(t -> 'b) * (t -> 'b -> unit) ->
   ?num_threads:int ->
   ?name:string ->
   'a
@@ -34,6 +35,9 @@ val create : (unit -> t, _) create_args
       [Domain.recommended_domain_count()], ie one worker per CPU core. On OCaml
       4 the default is [4] (since there is only one domain).
     @param on_exit_thread called at the end of each worker thread in the pool.
+    @param around_task
+      a pair of [before, after] functions ran around each task. See
+      {!Pool.create_args}.
     @param name name for the pool, used in tracing (since 0.6) *)
 
 val with_ : (unit -> (t -> 'a) -> 'a, _) create_args

src/core/hmap_ls_.real.ml

@@ -9,19 +9,19 @@ let k_local_hmap : Hmap.t FLS.t = FLS.create ()
 
 (** Access the local [hmap], or an empty one if not set *)
 let[@inline] get_local_hmap () : Hmap.t =
-  match TLS.get_exn k_cur_st with
+  match TLS.get_exn k_cur_fiber with
   | exception TLS.Not_set -> Hmap.empty
-  | { cur_fiber = fiber; _ } -> FLS.get fiber ~default:Hmap.empty k_local_hmap
+  | fiber -> FLS.get fiber ~default:Hmap.empty k_local_hmap
 
 let[@inline] set_local_hmap (h : Hmap.t) : unit =
-  match TLS.get_exn k_cur_st with
+  match TLS.get_exn k_cur_fiber with
   | exception TLS.Not_set -> ()
-  | { cur_fiber = fiber; _ } -> FLS.set fiber k_local_hmap h
+  | fiber -> FLS.set fiber k_local_hmap h
 
 let[@inline] update_local_hmap (f : Hmap.t -> Hmap.t) : unit =
-  match TLS.get_exn k_cur_st with
+  match TLS.get_exn k_cur_fiber with
   | exception TLS.Not_set -> ()
-  | { cur_fiber = fiber; _ } ->
+  | fiber ->
     let h = FLS.get fiber ~default:Hmap.empty k_local_hmap in
     let h = f h in
     FLS.set fiber k_local_hmap h

src/core/main.ml

@@ -1,7 +1,6 @@
 exception Oh_no of Exn_bt.t
 
 let main' ?(block_signals = false) () (f : Runner.t -> 'a) : 'a =
-  let module WL = Worker_loop_ in
   let worker_st =
     Fifo_pool.Private_.create_single_threaded_state ~thread:(Thread.self ())
       ~on_exn:(fun e bt -> raise (Oh_no (Exn_bt.make e bt)))
@@ -9,17 +8,15 @@ let main' ?(block_signals = false) () (f : Runner.t -> 'a) : 'a =
   in
   let runner = Fifo_pool.Private_.runner_of_state worker_st in
   try
-    let fut = Fut.spawn ~on:runner (fun () -> f runner) in
-    Fut.on_result fut (fun _ -> Runner.shutdown_without_waiting runner);
-
-    Thread_local_storage.set Runner.For_runner_implementors.k_cur_st
-      { cur_fiber = Picos.Fiber.create ~forbid:true fut; runner };
+    let fiber = Fut.spawn ~on:runner (fun () -> f runner) in
+    Fut.on_result fiber (fun _ -> Runner.shutdown_without_waiting runner);
 
     (* run the main thread *)
-    WL.worker_loop worker_st ~block_signals (* do not disturb existing thread *)
+    Worker_loop_.worker_loop worker_st
+      ~block_signals (* do not disturb existing thread *)
       ~ops:Fifo_pool.Private_.worker_ops;
 
-    match Fut.peek fut with
+    match Fut.peek fiber with
     | Some (Ok x) -> x
     | Some (Error ebt) -> Exn_bt.raise ebt
     | None -> assert false

src/core/runner.ml

@@ -47,12 +47,7 @@ module For_runner_implementors = struct
   let create ~size ~num_tasks ~shutdown ~run_async () : t =
     { size; num_tasks; shutdown; run_async }
 
-  type nonrec thread_local_state = thread_local_state = {
-    mutable runner: t;
-    mutable cur_fiber: fiber;
-  }
-
-  let k_cur_st : thread_local_state TLS.t = Types_.k_cur_st
+  let k_cur_runner : t TLS.t = Types_.k_cur_runner
 end
 
 let dummy : t =

src/core/runner.mli

@@ -72,13 +72,7 @@ module For_runner_implementors : sig
       {b NOTE}: the runner should support DLA and {!Suspend_} on OCaml 5.x, so
       that {!Fork_join} and other 5.x features work properly. *)
 
-  type thread_local_state = {
-    mutable runner: t;
-    mutable cur_fiber: fiber;
-  }
-  (** State set in thread-local-storage for worker threads *)
-
-  val k_cur_st : thread_local_state Thread_local_storage.t
+  val k_cur_runner : t Thread_local_storage.t
   (** 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. *)

src/core/types_.ml

@@ -11,12 +11,8 @@ type runner = {
   num_tasks: unit -> int;
 }
 
-type thread_local_state = {
-  mutable runner: runner;
-  mutable cur_fiber: fiber;
-}
-
-let k_cur_st : thread_local_state TLS.t = TLS.create ()
+let k_cur_runner : runner TLS.t = TLS.create ()
+let k_cur_fiber : fiber TLS.t = TLS.create ()
 
 let _dummy_computation : Picos.Computation.packed =
   let c = Picos.Computation.create () in
@@ -24,15 +20,11 @@ let _dummy_computation : Picos.Computation.packed =
   Picos.Computation.Packed c
 
 let _dummy_fiber = Picos.Fiber.create_packed ~forbid:true _dummy_computation
-
-let[@inline] get_current_runner () : _ option =
-  match TLS.get_exn k_cur_st with
-  | st -> Some st.runner
-  | exception TLS.Not_set -> None
+let[@inline] get_current_runner () : _ option = TLS.get_opt k_cur_runner
 
 let[@inline] get_current_fiber () : fiber option =
-  match TLS.get_exn k_cur_st with
-  | { cur_fiber = f; _ } when f != _dummy_fiber -> Some f
+  match TLS.get_exn k_cur_fiber with
+  | f when f != _dummy_fiber -> Some f
   | _ -> None
   | exception TLS.Not_set -> None
 
@@ -40,7 +32,7 @@ let error_get_current_fiber_ =
   "Moonpool: get_current_fiber was called outside of a fiber."
 
 let[@inline] get_current_fiber_exn () : fiber =
-  match TLS.get_exn k_cur_st with
-  | { cur_fiber = f; _ } when f != _dummy_fiber -> f
+  match TLS.get_exn k_cur_fiber with
+  | f when f != _dummy_fiber -> f
   | _ -> failwith error_get_current_fiber_
   | exception TLS.Not_set -> failwith error_get_current_fiber_

src/core/worker_loop_.ml

@@ -13,11 +13,15 @@ type task_full =
     }
       -> task_full
 
+type around_task =
+  | AT_pair : (Runner.t -> 'a) * (Runner.t -> 'a -> unit) -> around_task
+
 exception No_more_tasks
 
 type 'st ops = {
   schedule: 'st -> task_full -> unit;
   get_next_task: 'st -> task_full;  (** @raise No_more_tasks *)
+  around_task: 'st -> around_task;
   on_exn: 'st -> Exn_bt.t -> unit;
   runner: 'st -> Runner.t;
   before_start: 'st -> unit;
@@ -102,13 +106,7 @@ end
 module Fine_grained (Args : FINE_GRAINED_ARGS) () = struct
   open Args
 
-  let cur_st : Runner.For_runner_implementors.thread_local_state Lazy.t =
-    lazy
-      (match TLS.get_exn Runner.For_runner_implementors.k_cur_st with
-      | st -> st
-      | exception TLS.Not_set ->
-        failwith "Moonpool: worker loop: no current state set")
-
+  let cur_fiber : fiber ref = ref _dummy_fiber
   let runner = ops.runner st
 
   type state =
@@ -119,12 +117,15 @@ module Fine_grained (Args : FINE_GRAINED_ARGS) () = struct
   let state = ref New
 
   let run_task (task : task_full) : unit =
+    let (AT_pair (before_task, after_task)) = ops.around_task st in
     let fiber =
       match task with
       | T_start { fiber; _ } | T_resume { fiber; _ } -> fiber
     in
 
-    (Lazy.force cur_st).cur_fiber <- fiber;
+    cur_fiber := fiber;
+    TLS.set k_cur_fiber fiber;
+    let _ctx = before_task runner in
 
     (* run the task now, catching errors, handling effects *)
     assert (task != _dummy_task);
@@ -139,7 +140,10 @@ module Fine_grained (Args : FINE_GRAINED_ARGS) () = struct
        let ebt = Exn_bt.make e bt in
        ops.on_exn st ebt);
 
-    (Lazy.force cur_st).cur_fiber <- _dummy_fiber
+    after_task runner _ctx;
+
+    cur_fiber := _dummy_fiber;
+    TLS.set k_cur_fiber _dummy_fiber
 
   let setup ~block_signals () : unit =
     if !state <> New then invalid_arg "worker_loop.setup: not a new instance";
@@ -162,9 +166,9 @@ module Fine_grained (Args : FINE_GRAINED_ARGS) () = struct
       with _ -> ()
     );
 
-    ops.before_start st;
-    (Lazy.force cur_st).runner <- runner;
-    ()
+    TLS.set Runner.For_runner_implementors.k_cur_runner runner;
+
+    ops.before_start st
 
   let run ?(max_tasks = max_int) () : unit =
     if !state <> Ready then invalid_arg "worker_loop.run: not setup";
@@ -182,7 +186,7 @@ module Fine_grained (Args : FINE_GRAINED_ARGS) () = struct
   let teardown () =
     if !state <> Torn_down then (
       state := Torn_down;
-      (Lazy.force cur_st).cur_fiber <- _dummy_fiber;
+      cur_fiber := _dummy_fiber;
       ops.cleanup st
     )
 end

src/core/worker_loop_.mli

@@ -18,11 +18,15 @@ type task_full =
 
 val _dummy_task : task_full
 
+type around_task =
+  | AT_pair : (Runner.t -> 'a) * (Runner.t -> 'a -> unit) -> around_task
+
 exception No_more_tasks
 
 type 'st ops = {
   schedule: 'st -> task_full -> unit;
   get_next_task: 'st -> task_full;
+  around_task: 'st -> around_task;
   on_exn: 'st -> Exn_bt.t -> unit;
   runner: 'st -> Runner.t;
   before_start: 'st -> unit;

src/core/ws_pool.ml

@@ -28,6 +28,7 @@ type state = {
   cond: Condition.t;
   mutable as_runner: t;
   (* init options *)
+  around_task: WL.around_task;
   name: string option;
   on_init_thread: dom_id:int -> t_id:int -> unit -> unit;
   on_exit_thread: dom_id:int -> t_id:int -> unit -> unit;
@@ -55,8 +56,7 @@ let num_tasks_ (self : state) : int =
   !n
 
 (** TLS, used by worker to store their specific state and be able to retrieve it
-    from tasks when we schedule new sub-tasks. This way we can schedule the new
-    task directly in the local work queue, where it might be stolen. *)
+    from tasks when we schedule new sub-tasks. *)
 let k_worker_state : worker_state TLS.t = TLS.create ()
 
 let[@inline] get_current_worker_ () : worker_state option =
@@ -180,8 +180,8 @@ and wait_on_main_queue (self : worker_state) : WL.task_full =
 let before_start (self : worker_state) : unit =
   let t_id = Thread.id @@ Thread.self () in
   self.st.on_init_thread ~dom_id:self.dom_id ~t_id ();
-  TLS.set Runner.For_runner_implementors.k_cur_st
-    { cur_fiber = _dummy_fiber; runner = self.st.as_runner };
+  TLS.set k_cur_fiber _dummy_fiber;
+  TLS.set Runner.For_runner_implementors.k_cur_runner self.st.as_runner;
   TLS.set k_worker_state self;
 
   (* set thread name *)
@@ -198,6 +198,7 @@ let cleanup (self : worker_state) : unit =
 
 let worker_ops : worker_state WL.ops =
   let runner (st : worker_state) = st.st.as_runner in
+  let around_task st = st.st.around_task in
   let on_exn (st : worker_state) (ebt : Exn_bt.t) =
     st.st.on_exn (Exn_bt.exn ebt) (Exn_bt.bt ebt)
   in
@@ -205,6 +206,7 @@ let worker_ops : worker_state WL.ops =
     WL.schedule = schedule_from_w;
     runner;
     get_next_task;
+    around_task;
     on_exn;
     before_start;
     cleanup;
@@ -233,6 +235,7 @@ type ('a, 'b) create_args =
   ?on_init_thread:(dom_id:int -> t_id:int -> unit -> unit) ->
   ?on_exit_thread:(dom_id:int -> t_id:int -> unit -> unit) ->
   ?on_exn:(exn -> Printexc.raw_backtrace -> unit) ->
+  ?around_task:(t -> 'b) * (t -> 'b -> unit) ->
   ?num_threads:int ->
   ?name:string ->
   'a
@@ -240,8 +243,15 @@ type ('a, 'b) create_args =
 
 let create ?(on_init_thread = default_thread_init_exit_)
     ?(on_exit_thread = default_thread_init_exit_) ?(on_exn = fun _ _ -> ())
-    ?num_threads ?name () : t =
+    ?around_task ?num_threads ?name () : t =
   let pool_id_ = Id.create () in
+  (* wrapper *)
+  let around_task =
+    match around_task with
+    | Some (f, g) -> WL.AT_pair (f, g)
+    | None -> WL.AT_pair (ignore, fun _ _ -> ())
+  in
+
   let num_domains = Domain_pool_.max_number_of_domains () in
   let num_threads = Util_pool_.num_threads ?num_threads () in
 
@@ -258,6 +268,7 @@ let create ?(on_init_thread = default_thread_init_exit_)
       n_waiting_nonzero = true;
       mutex = Mutex.create ();
       cond = Condition.create ();
+      around_task;
       on_exn;
       on_init_thread;
       on_exit_thread;
@@ -313,9 +324,11 @@ let create ?(on_init_thread = default_thread_init_exit_)
 
   pool.as_runner
 
-let with_ ?on_init_thread ?on_exit_thread ?on_exn ?num_threads ?name () f =
+let with_ ?on_init_thread ?on_exit_thread ?on_exn ?around_task ?num_threads
+    ?name () f =
   let pool =
-    create ?on_init_thread ?on_exit_thread ?on_exn ?num_threads ?name ()
+    create ?on_init_thread ?on_exit_thread ?on_exn ?around_task ?num_threads
+      ?name ()
   in
   let@ () = Fun.protect ~finally:(fun () -> shutdown pool) in
   f pool

src/core/ws_pool.mli

@@ -24,6 +24,7 @@ type ('a, 'b) create_args =
   ?on_init_thread:(dom_id:int -> t_id:int -> unit -> unit) ->
   ?on_exit_thread:(dom_id:int -> t_id:int -> unit -> unit) ->
   ?on_exn:(exn -> Printexc.raw_backtrace -> unit) ->
+  ?around_task:(t -> 'b) * (t -> 'b -> unit) ->
   ?num_threads:int ->
   ?name:string ->
   'a
@@ -39,6 +40,11 @@ val create : (unit -> t, _) create_args
       [Domain.recommended_domain_count()], ie one worker thread per CPU core. On
       OCaml 4 the default is [4] (since there is only one domain).
     @param on_exit_thread called at the end of each thread in the pool
+    @param around_task
+      a pair of [before, after], where [before pool] is called before a task is
+      processed, on the worker thread about to run it, and returns [x]; and
+      [after pool x] is called by the same thread after the task is over. (since
+      0.2)
     @param name
       a name for this thread pool, used if tracing is enabled (since 0.6) *)
 

src/dpool/moonpool_dpool.ml

@@ -146,7 +146,7 @@ let work_ idx (st : worker_state) : unit =
 let () =
   assert (Domain_.is_main_domain ());
   let w = { th_count = Atomic.make 1; q = Bb_queue.create () } in
-  (* thread that stays alive since [th_count>0] will always hold *)
+  (* thread that stays alive *)
   ignore (Thread.create (fun () -> work_ 0 w) () : Thread.t);
   domains_.(0) <- Lock.create (Some w, None)
 
@@ -154,8 +154,7 @@ let[@inline] max_number_of_domains () : int = Array.length domains_
 
 let run_on (i : int) (f : unit -> unit) : unit =
   assert (i < Array.length domains_);
-
-  let w : worker_state =
+  let w =
     Lock.update_map domains_.(i) (function
       | (Some w, _) as st ->
         Atomic.incr w.th_count;

src/lwt/moonpool_lwt.ml

@@ -7,6 +7,8 @@ end
 
 module Fut = Moonpool.Fut
 
+let default_around_task_ : WL.around_task = AT_pair (ignore, fun _ _ -> ())
+
 let on_uncaught_exn : (Moonpool.Exn_bt.t -> unit) ref =
   ref (fun ebt ->
       Printf.eprintf "uncaught exception in moonpool-lwt:\n%s" (Exn_bt.show ebt))
@@ -88,6 +90,8 @@ end
 module Ops = struct
   type st = Scheduler_state.st
 
+  let around_task _ = default_around_task_
+
   let schedule (self : st) t =
     if Atomic.get self.closed then
       failwith "moonpool-lwt.schedule: scheduler is closed";
@@ -118,7 +122,15 @@ module Ops = struct
     ()
 
   let ops : st WL.ops =
-    { schedule; get_next_task; on_exn; runner; before_start; cleanup }
+    {
+      schedule;
+      around_task;
+      get_next_task;
+      on_exn;
+      runner;
+      before_start;
+      cleanup;
+    }
 
   let setup st =
     if Atomic.compare_and_set Scheduler_state.cur_st None (Some st) then

test/t_unfair.ml

@@ -14,11 +14,14 @@ let run ~kind () =
   let pool =
     let on_init_thread ~dom_id:_ ~t_id () =
       Trace.set_thread_name (Printf.sprintf "pool worker %d" t_id)
+    and around_task =
+      ( (fun self -> Trace.counter_int "n_tasks" (Ws_pool.num_tasks self)),
+        fun self () -> Trace.counter_int "n_tasks" (Ws_pool.num_tasks self) )
     in
 
     match kind with
-    | `Simple -> Fifo_pool.create ~num_threads:3 ~on_init_thread ()
-    | `Ws_pool -> Ws_pool.create ~num_threads:3 ~on_init_thread ()
+    | `Simple -> Fifo_pool.create ~num_threads:3 ~on_init_thread ~around_task ()
+    | `Ws_pool -> Ws_pool.create ~num_threads:3 ~on_init_thread ~around_task ()
   in
 
   (* make all threads busy *)