From b6db37fab00d91b9b16cea8e79ea876a8574d088 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Fri, 9 Jun 2023 16:07:49 -0400
Subject: [PATCH] add thread/process name setting

---
 src/collector.ml     |  6 ++++++
 src/tef/trace_tef.ml | 30 ++++++++++++++++++++++++++++++
 src/trace.ml         | 10 ++++++++++
 src/trace.mli        |  6 ++++++
 4 files changed, 52 insertions(+)

diff --git a/src/collector.ml b/src/collector.ml
index fef5edd..1b45212 100644
--- a/src/collector.ml
+++ b/src/collector.ml
@@ -28,6 +28,12 @@ module type S = sig
     string ->
     unit
 
+  val name_thread : string -> unit
+  (** Give a name to the current thread *)
+
+  val name_process : string -> unit
+  (** Give a name to the current process *)
+
   val shutdown : unit -> unit
   (** Shutdown collector, possibly waiting for it to finish sending data. *)
 end
diff --git a/src/tef/trace_tef.ml b/src/tef/trace_tef.ml
index 18718c5..59d37b5 100644
--- a/src/tef/trace_tef.ml
+++ b/src/tef/trace_tef.ml
@@ -60,6 +60,11 @@ type event =
       id: span;
       time_us: float;
     }
+  | E_name_process of { name: string }
+  | E_name_thread of {
+      tid: int;
+      name: string;
+    }
 
 module Span_tbl = Hashtbl.Make (struct
   include Int64
@@ -172,6 +177,23 @@ module Writer = struct
       (emit_args_o_ pp_user_data_)
       args;
     ()
+
+  let emit_name_thread ~tid ~name (self : t) : unit =
+    emit_sep_ self;
+    Printf.fprintf self.oc
+      {json|{"pid": %d,"tid": %d,"name":"thread_name","ph":"M"%a}|json} self.pid
+      tid
+      (emit_args_o_ pp_user_data_)
+      [ "name", `String name ];
+    ()
+
+  let emit_name_process ~name (self : t) : unit =
+    emit_sep_ self;
+    Printf.fprintf self.oc
+      {json|{"pid": %d,"name":"process_name","ph":"M"%a}|json} self.pid
+      (emit_args_o_ pp_user_data_)
+      [ "name", `String name ];
+    ()
 end
 
 let bg_thread ~out (events : event B_queue.t) : unit =
@@ -205,6 +227,8 @@ let bg_thread ~out (events : event B_queue.t) : unit =
         Span_tbl.remove spans id;
         Writer.emit_duration_event ~tid ~name ~start:start_us ~end_:stop_us
           ~args:data writer)
+    | E_name_process { name } -> Writer.emit_name_process ~name writer
+    | E_name_thread { tid; name } -> Writer.emit_name_thread ~tid ~name writer
   in
 
   try
@@ -275,6 +299,12 @@ let collector ~out () : collector =
       B_queue.push events
         (E_message
            { (* __FUNCTION__; __FILE__; __LINE__; *) tid; time_us; msg; data })
+
+    let name_process name : unit = B_queue.push events (E_name_process { name })
+
+    let name_thread name : unit =
+      let tid = get_tid_ () in
+      B_queue.push events (E_name_thread { tid; name })
   end in
   (module M)
 
diff --git a/src/trace.ml b/src/trace.ml
index 17f5502..4843e21 100644
--- a/src/trace.ml
+++ b/src/trace.ml
@@ -76,6 +76,16 @@ let messagef ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data k =
             C.message ?__FUNCTION__ ~__FILE__ ~__LINE__ ~data str)
           fmt)
 
+let set_thread_name name : unit =
+  match A.get collector with
+  | None -> ()
+  | Some (module C) -> C.name_thread name
+
+let set_process_name name : unit =
+  match A.get collector with
+  | None -> ()
+  | Some (module C) -> C.name_process name
+
 let setup_collector c : unit =
   while
     let cur = A.get collector in
diff --git a/src/trace.mli b/src/trace.mli
index 1c741d3..5222614 100644
--- a/src/trace.mli
+++ b/src/trace.mli
@@ -48,6 +48,12 @@ val messagef :
   ((('a, Format.formatter, unit, unit) format4 -> 'a) -> unit) ->
   unit
 
+val set_thread_name : string -> unit
+(** Give a name to the current thread. *)
+
+val set_process_name : string -> unit
+(** Give a name to the current process. *)
+
 (** {2 Collector} *)
 
 type collector = (module Collector.S)