Merge pull request #14811 from MinaProtocol/fix/trace-tool

buildkite/src/Jobs/Release/TraceTool.dhall

@@ -18,19 +18,19 @@ in
 Pipeline.build
   Pipeline.Config::{
     spec = JobSpec::{
-      dirtyWhen = [ S.contains "src/app/trace-tool", S.strictlyStart (S.contains "buildkite/src/Jobs/TraceTool") ],
+      dirtyWhen = [ S.contains "src/app/trace-tool", S.strictlyStart (S.contains "buildkite/src/Jobs/Release/TraceTool") ],
       path = "Release",
       name = "TraceTool"
     },
     steps = [
       Command.build
         Command.Config::{
-          commands = RunInToolchain.runInToolchain ([] : List Text) "cd src/app/trace-tool && PATH=/home/opam/.cargo/bin:$PATH cargo build"
+          commands = RunInToolchain.runInToolchain ([] : List Text) "cd src/app/trace-tool && PATH=/home/opam/.cargo/bin:$PATH cargo build --release"
           , label = "Build trace-tool"
           , key = "build-trace-tool"
           , target = Size.Small
           , docker = None Docker.Type
-          , artifact_paths = [ S.contains "src/app/trace-tool/target/debug/trace-tool" ]
+          , artifact_paths = [ S.contains "src/app/trace-tool/target/release/trace-tool" ]
         }
     ]
   }

src/lib/o1trace/execution_timer.ml

@@ -26,4 +26,8 @@ let on_job_enter _fiber = ()
 
 let on_job_exit fiber elapsed_time = record_elapsed_time fiber elapsed_time
 
+let on_new_fiber _fiber = ()
+
+let on_cycle_end () = ()
+
 let elapsed_time_of_thread thread = !(Plugins.plugin_state (module T) thread)

src/lib/o1trace/o1trace.ml

@@ -22,6 +22,10 @@ let on_job_exit ctx elapsed_time =
   Option.iter (Thread.Fiber.of_context ctx) ~f:(fun thread ->
       on_job_exit' thread elapsed_time )
 
+let on_new_fiber (fiber : Thread.Fiber.t) =
+  Plugins.dispatch (fun (module Plugin : Plugins.Plugin_intf) ->
+      Plugin.on_new_fiber fiber )
+
 let current_sync_fiber = ref None
 
 (* grabs the parent fiber, returning the fiber (if available) and a reset function to call after exiting the child fiber *)
@@ -72,7 +76,8 @@ let exec_thread ~exec_same_thread ~exec_new_thread name =
         | Some fiber ->
             fiber
         | None ->
-            Thread.Fiber.register name parent
+            let fib = Thread.Fiber.register name parent in
+            on_new_fiber fib ; fib
       in
       exec_new_thread fiber
   in
@@ -113,7 +118,11 @@ let sync_thread name f =
           on_job_exit' fiber elapsed_time ;
           result )
 
-let () = Stdlib.(Async_kernel.Tracing.fns := { on_job_enter; on_job_exit })
+let () =
+  Stdlib.(Async_kernel.Tracing.fns := { on_job_enter; on_job_exit }) ;
+  Scheduler.Expert.run_every_cycle_end (fun () ->
+      Plugins.dispatch (fun (module Plugin : Plugins.Plugin_intf) ->
+          Plugin.on_cycle_end () ) )
 
 (*
 let () =

src/lib/o1trace/o1trace.mli

@@ -15,8 +15,13 @@ module Thread : sig
   val dump_thread_graph : unit -> bytes
 
   module Fiber : sig
-    type t = Thread.Fiber.t = { id : int; parent : t option; thread : Thread.t }
+    type t = Thread.Fiber.t =
+      { id : int; parent : t option; thread : Thread.t; key : string list }
+
+    val key : t -> string list
   end
+
+  val iter_fibers : f:(Fiber.t -> unit) -> unit
 end
 
 module Plugins : module type of Plugins

src/lib/o1trace/plugins.ml

@@ -29,6 +29,10 @@ module type Plugin_intf = sig
   val on_job_enter : Thread.Fiber.t -> unit
 
   val on_job_exit : Thread.Fiber.t -> Time_ns.Span.t -> unit
+
+  val on_new_fiber : Thread.Fiber.t -> unit
+
+  val on_cycle_end : unit -> unit
 end
 
 module Register_plugin (Plugin_spec : Plugin_spec_intf) () :

src/lib/o1trace/thread.ml

@@ -69,7 +69,8 @@ module Fiber = struct
 
   let next_id = ref 1
 
-  type t = { id : int; parent : t option; thread : thread } [@@deriving sexp_of]
+  type t = { id : int; parent : t option; thread : thread; key : string list }
+  [@@deriving sexp_of]
 
   let ctx_id : t Type_equal.Id.t = Type_equal.Id.create ~name:"fiber" sexp_of_t
 
@@ -87,7 +88,7 @@ module Fiber = struct
         fiber
     | None ->
         let thread = register name in
-        let fiber = { id = !next_id; parent; thread } in
+        let fiber = { id = !next_id; parent; thread; key } in
         incr next_id ;
         Hashtbl.set fibers ~key ~data:fiber ;
         Option.iter parent ~f:(fun p -> Graph.add_edge graph p.thread.name name) ;
@@ -98,8 +99,12 @@ module Fiber = struct
     Execution_context.with_local ctx ctx_id (Some t)
 
   let of_context ctx = Execution_context.find_local ctx ctx_id
+
+  let key { thread = { name; _ }; parent; _ } = fiber_key name parent
 end
 
 let of_context ctx =
   let%map.Option fiber = Fiber.of_context ctx in
   fiber.thread
+
+let iter_fibers ~f = Hashtbl.iter Fiber.fibers ~f

src/lib/o1trace/webkit_event/o1trace_webkit_event.ml

@@ -5,14 +5,18 @@ module Scheduler = Async_kernel_scheduler
 
 let current_wr = ref None
 
-let emit_event =
+let emitted_since_cycle_ended = ref false
+
+let emit_event' =
   let buf = Bigstring.create 512 in
-  fun event ->
-    Option.iter !current_wr ~f:(fun wr ->
-        try Webkit_trace_event_binary_output.emit_event ~buf wr event
-        with exn ->
-          Writer.writef wr "failed to write o1trace event: %s\n"
-            (Exn.to_string exn) )
+  fun wr event ->
+    emitted_since_cycle_ended := true ;
+    try Webkit_trace_event_binary_output.emit_event ~buf wr event
+    with exn ->
+      Writer.writef wr "failed to write o1trace event: %s\n" (Exn.to_string exn)
+
+let emit_event event =
+  Option.iter !current_wr ~f:(fun wr -> emit_event' wr event)
 
 let timestamp () =
   Time_stamp_counter.now () |> Time_stamp_counter.to_int63 |> Int63.to_int_exn
@@ -28,22 +32,8 @@ let new_event (k : event_kind) : event =
   ; tid = 0
   }
 
-(* This will track ids per thread. If we need to track ids per fiber,
-   we will need to feed the fiber id into the plugin hooks. *)
-let id_of_thread =
-  let ids = String.Table.create () in
-  let next_id = ref 0 in
-  let alloc_id () =
-    let id = !next_id in
-    incr next_id ; id
-  in
-  fun thread_name -> Hashtbl.find_or_add ids thread_name ~default:alloc_id
-
-let new_thread_event ?(include_name = false) thread_name event_kind =
-  { (new_event event_kind) with
-    tid = id_of_thread thread_name
-  ; name = (if include_name then thread_name else "")
-  }
+let new_thread_event ?(include_name = "") tid event_kind =
+  { (new_event event_kind) with tid; name = include_name }
 
 (*
 
@@ -97,17 +87,22 @@ module T = struct
       end)
       ()
 
+  let most_recent_id = ref 0
+
   let on_job_enter (fiber : O1trace.Thread.Fiber.t) =
-    emit_event
-      (new_thread_event (O1trace.Thread.name fiber.thread) Thread_switch)
+    if fiber.id <> !most_recent_id then (
+      most_recent_id := fiber.id ;
+      emit_event (new_thread_event fiber.id Thread_switch) )
 
   let on_job_exit _fiber _time_elapsed = ()
 
-  (*
+  let on_new_fiber (fiber : O1trace.Thread.Fiber.t) =
+    let fullname = String.concat ~sep:"/" (O1trace.Thread.Fiber.key fiber) in
+    emit_event (new_thread_event ~include_name:fullname fiber.id New_thread)
+
   let on_cycle_end () =
-    let sch = Scheduler.t () in
-    emit_event (new_thread_event thread_name Cycle_end) ;
-  *)
+    if !emitted_since_cycle_ended then emit_event (new_event Cycle_end) ;
+    emitted_since_cycle_ended := false
 end
 
 let start_tracing wr =
@@ -116,11 +111,7 @@ let start_tracing wr =
   else (
     current_wr := Some wr ;
     emit_event (new_event Pid_is) ;
-    O1trace.Thread.iter_threads ~f:(fun thread ->
-        emit_event
-          (new_thread_event ~include_name:true
-             (O1trace.Thread.name thread)
-             New_thread ) ) ;
+    O1trace.Thread.iter_fibers ~f:T.on_new_fiber ;
     O1trace.Plugins.enable_plugin (module T) )
 
 let stop_tracing () =