From 27b213e30f3641f4d9c291a71904ffb36d66883c Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Wed, 7 Feb 2024 13:31:48 -0500
Subject: [PATCH] tracing: add optional name for pools

---
 src/fifo_pool.ml      | 9 ++++++++-
 src/fifo_pool.mli     | 2 ++
 src/tracing_.dummy.ml | 1 +
 src/tracing_.mli      | 1 +
 src/tracing_.real.ml  | 1 +
 src/ws_pool.ml        | 9 ++++++++-
 src/ws_pool.mli       | 2 ++
 7 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/src/fifo_pool.ml b/src/fifo_pool.ml
index c3a09e02..c4253c32 100644
--- a/src/fifo_pool.ml
+++ b/src/fifo_pool.ml
@@ -78,11 +78,12 @@ type ('a, 'b) create_args =
   ?on_exn:(exn -> Printexc.raw_backtrace -> unit) ->
   ?around_task:(t -> 'b) * (t -> 'b -> unit) ->
   ?num_threads:int ->
+  ?name:string ->
   'a
 
 let create ?(on_init_thread = default_thread_init_exit_)
     ?(on_exit_thread = default_thread_init_exit_) ?(on_exn = fun _ _ -> ())
-    ?around_task ?num_threads () : t =
+    ?around_task ?num_threads ?name () : t =
   (* wrapper *)
   let around_task =
     match around_task with
@@ -128,6 +129,12 @@ let create ?(on_init_thread = default_thread_init_exit_)
       let t_id = Thread.id thread in
       on_init_thread ~dom_id:dom_idx ~t_id ();
 
+      (* set thread name *)
+      Option.iter
+        (fun name ->
+          Tracing_.set_thread_name (Printf.sprintf "%s.worker.%d" name i))
+        name;
+
       let run () = worker_thread_ pool runner ~on_exn ~around_task in
 
       (* now run the main loop *)
diff --git a/src/fifo_pool.mli b/src/fifo_pool.mli
index 9b82081b..25cde495 100644
--- a/src/fifo_pool.mli
+++ b/src/fifo_pool.mli
@@ -22,6 +22,7 @@ type ('a, 'b) create_args =
   ?on_exn:(exn -> Printexc.raw_backtrace -> unit) ->
   ?around_task:(t -> 'b) * (t -> 'b -> unit) ->
   ?num_threads:int ->
+  ?name:string ->
   'a
 (** Arguments used in {!create}. See {!create} for explanations. *)
 
@@ -35,6 +36,7 @@ val create : (unit -> t, _) create_args
      @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 NEXT_RELEASE)
   *)
 
 val with_ : (unit -> (t -> 'a) -> 'a, _) create_args
diff --git a/src/tracing_.dummy.ml b/src/tracing_.dummy.ml
index 9977c6a8..d685a5b3 100644
--- a/src/tracing_.dummy.ml
+++ b/src/tracing_.dummy.ml
@@ -2,3 +2,4 @@ let enabled () = false
 let dummy_span = 0L
 let enter_span _name = dummy_span
 let exit_span = ignore
+let set_thread_name = ignore
diff --git a/src/tracing_.mli b/src/tracing_.mli
index beadc19f..35379332 100644
--- a/src/tracing_.mli
+++ b/src/tracing_.mli
@@ -2,3 +2,4 @@ val dummy_span : int64
 val enter_span : string -> int64
 val exit_span : int64 -> unit
 val enabled : unit -> bool
+val set_thread_name : string -> unit
diff --git a/src/tracing_.real.ml b/src/tracing_.real.ml
index a71d9b3b..f71ec418 100644
--- a/src/tracing_.real.ml
+++ b/src/tracing_.real.ml
@@ -3,6 +3,7 @@ module Trace = Trace_core
 let enabled = Trace.enabled
 let dummy_span = Int64.min_int
 let dummy_file_ = "<unknown file>"
+let set_thread_name = Trace.set_thread_name
 
 let[@inline] enter_span name : int64 =
   if name = "" then
diff --git a/src/ws_pool.ml b/src/ws_pool.ml
index 7fa930d1..cb87eeba 100644
--- a/src/ws_pool.ml
+++ b/src/ws_pool.ml
@@ -247,6 +247,7 @@ type ('a, 'b) create_args =
   ?on_exn:(exn -> Printexc.raw_backtrace -> unit) ->
   ?around_task:(t -> 'b) * (t -> 'b -> unit) ->
   ?num_threads:int ->
+  ?name:string ->
   'a
 (** Arguments used in {!create}. See {!create} for explanations. *)
 
@@ -254,7 +255,7 @@ 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 _ _ -> ())
-    ?around_task ?num_threads () : t =
+    ?around_task ?num_threads ?name () : t =
   let pool_id_ = Id.create () in
   (* wrapper *)
   let around_task =
@@ -320,6 +321,12 @@ let create ?(on_init_thread = default_thread_init_exit_)
       let t_id = Thread.id thread in
       on_init_thread ~dom_id:dom_idx ~t_id ();
 
+      (* set thread name *)
+      Option.iter
+        (fun name ->
+          Tracing_.set_thread_name (Printf.sprintf "%s.worker.%d" name i))
+        name;
+
       let run () = worker_thread_ pool ~runner w in
 
       (* now run the main loop *)
diff --git a/src/ws_pool.mli b/src/ws_pool.mli
index 578966a3..50a79fe3 100644
--- a/src/ws_pool.mli
+++ b/src/ws_pool.mli
@@ -27,6 +27,7 @@ type ('a, 'b) create_args =
   ?on_exn:(exn -> Printexc.raw_backtrace -> unit) ->
   ?around_task:(t -> 'b) * (t -> 'b -> unit) ->
   ?num_threads:int ->
+  ?name:string ->
   'a
 (** Arguments used in {!create}. See {!create} for explanations. *)
 
@@ -44,6 +45,7 @@ val create : (unit -> t, _) create_args
       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 NEXT_RELEASE)
   *)
 
 val with_ : (unit -> (t -> 'a) -> 'a, _) create_args