more tracing, use local store key for trace nameing

raphael-proust · raphael-proust · commit de6715bc660d · 2025-10-18T15:52:26.000+08:00
diff --git a/src/core/lwt.ml b/src/core/lwt.ml
@@ -1671,6 +1671,10 @@ include Pending_promises
 
 
 
+let tracing_context = new_key ()
+let tracing_counter = ref 0
+let tracing_counter () = incr tracing_counter; !tracing_counter
+
 module Sequential_composition :
 sig
   (* Main interface (public) *)
@@ -1903,6 +1907,8 @@ struct
       p''
 
   let backtrace_bind (loc_f, loc_l, _, _) add_loc p f =
+    let trace_context = Sequence_associated_storage.get tracing_context in
+    let trace_counter = tracing_counter () in
     let Internal p = to_internal_promise p in
     let p = underlying p in
 
@@ -1912,7 +1918,7 @@ struct
       let saved_storage = !current_storage in
 
       let callback p_result =
-        Lwt_rte.emit_trace End loc_f loc_l;
+        Lwt_rte.emit_trace End trace_context trace_counter loc_f loc_l;
         match p_result with
         | Fulfilled v ->
           current_storage := saved_storage;
@@ -1956,7 +1962,7 @@ struct
       to_public_promise {state = Rejected (add_loc exn)}
 
     | Pending p_callbacks ->
-      Lwt_rte.emit_trace Begin loc_f loc_l;
+      Lwt_rte.emit_trace Begin trace_context trace_counter loc_f loc_l;
       let (p'', callback) = create_result_promise_and_callback_if_deferred () in
       add_implicitly_removed_callback p_callbacks callback;
       p''
@@ -2085,6 +2091,8 @@ struct
       p''
 
   let backtrace_catch (loc_f, loc_l, _, _) add_loc f h =
+    let trace_context = Sequence_associated_storage.get tracing_context in
+    let trace_counter = tracing_counter () in
     let p =
       try f ()
       with exn when Exception_filter.run exn -> fail exn
@@ -2098,7 +2106,7 @@ struct
       let saved_storage = !current_storage in
 
       let callback p_result =
-        Lwt_rte.emit_trace End loc_f loc_l;
+        Lwt_rte.emit_trace End trace_context trace_counter loc_f loc_l;
         match p_result with
         | Fulfilled _ as p_result ->
           let State_may_now_be_pending_proxy p'' = may_now_be_proxy p'' in
@@ -2143,7 +2151,7 @@ struct
           (p'', callback, p.state))
 
     | Pending p_callbacks ->
-      Lwt_rte.emit_trace Begin loc_f loc_l;
+      Lwt_rte.emit_trace Begin trace_context trace_counter loc_f loc_l;
       let (p'', callback) = create_result_promise_and_callback_if_deferred () in
       add_implicitly_removed_callback p_callbacks callback;
       p''
@@ -2224,6 +2232,8 @@ struct
       p''
 
   let backtrace_try_bind (loc_f, loc_l, _, _) add_loc f f' h =
+    let trace_context = Sequence_associated_storage.get tracing_context in
+    let trace_counter = tracing_counter () in
     let p =
       try f ()
       with exn when Exception_filter.run exn -> fail exn
@@ -2237,7 +2247,7 @@ struct
       let saved_storage = !current_storage in
 
       let callback p_result =
-        Lwt_rte.emit_trace End loc_f loc_l;
+        Lwt_rte.emit_trace End trace_context trace_counter loc_f loc_l;
         match p_result with
         | Fulfilled v ->
           current_storage := saved_storage;
@@ -2297,7 +2307,7 @@ struct
           (p'', callback, p.state))
 
     | Pending p_callbacks ->
-      Lwt_rte.emit_trace Begin loc_f loc_l;
+      Lwt_rte.emit_trace Begin trace_context trace_counter loc_f loc_l;
       let (p'', callback) = create_result_promise_and_callback_if_deferred () in
       add_implicitly_removed_callback p_callbacks callback;
       p''
diff --git a/src/core/lwt.mli b/src/core/lwt.mli
@@ -2076,3 +2076,5 @@ module Private : sig
     val current_storage : storage ref
   end
 end [@@alert trespassing "for internal use only, keep away"]
+
+val tracing_context : string key
diff --git a/src/core/lwt_rte.mli b/src/core/lwt_rte.mli
@@ -4,4 +4,4 @@ val emit_sch_call_end : unit -> unit
 val emit_sch_lap : unit -> unit
 val emit_job_count : int -> unit
 type span = Begin | End
-val emit_trace : span -> string -> int -> unit
+val emit_trace : span -> string option -> int -> string -> int -> unit
diff --git a/src/core/lwt_rte.with.ml b/src/core/lwt_rte.with.ml
@@ -4,6 +4,4 @@ let emit_sch_call_end () = Lwt_runtime_events.emit_sch_call_end ()
 let emit_sch_lap v = Lwt_runtime_events.emit_sch_lap v
 let emit_job_count v = Lwt_runtime_events.emit_job_count v
 type span = Runtime_events.Type.span = Begin | End
-let emit_trace k f l =
-  let s = Printf.sprintf "%s:%d" f l in
-  Lwt_runtime_events.Trace.emit_span (k, s)
+let emit_trace kind context count filename line = Lwt_runtime_events.Trace.emit {kind; context; count; filename; line}
diff --git a/src/core/lwt_rte.without.ml b/src/core/lwt_rte.without.ml
@@ -4,4 +4,4 @@ let emit_sch_call_end () = ()
 let emit_sch_lap _v = ()
 let emit_job_count _v = ()
 type span = Begin | End
-let emit_trace _k _f _l = ()
+let emit_trace _ _ _ _ _ = ()
diff --git a/src/runtime_events/lwt_runtime_events.ml b/src/runtime_events/lwt_runtime_events.ml
@@ -17,29 +17,65 @@ let emit_job_count v = Runtime_events.User.write unix_job_count v
 
 module Trace = struct
 
-  type labelled_span = Runtime_events.Type.span * string
-  let labelled_span : labelled_span Runtime_events.Type.t =
-    Runtime_events.Type.register
-      ~encode:(fun b (k, s) ->
-        let () = match k with
-          | Runtime_events.Type.Begin -> Bytes.set b 0 'B'
-          | End -> Bytes.set b 0 'E'
-        in
-        let l = min (String.length s) (Bytes.length b - 1) in
-        Bytes.blit_string s 0 b 1 l;
-        (l + 1))
-      ~decode:(fun b i ->
-        if i < 1 then failwith "unreadable tag for labelled_span";
-        let k = match Bytes.get b 0 with
-          | 'B' -> Runtime_events.Type.Begin
-          | 'E' -> End
-          | _ -> failwith "unreadable tag for labelled_span";
-        in
-        let s = Bytes.sub_string b 1 (i - 1) in
-        (k, s))
-
-  type Runtime_events.User.tag += LabelledSpan
-  let span = Runtime_events.User.register "lwt-trace" LabelledSpan labelled_span
-  let emit_span labelled_span = Runtime_events.User.write span labelled_span
+  type t =
+    { kind: Runtime_events.Type.span;
+      context: string option;
+      count: int;
+      filename: string;
+      line: int; }
+
+  let decode b i =
+    let offset = 0 in
+    let kind = match Bytes.get b offset with
+      | 'B' -> Runtime_events.Type.Begin
+      | 'E' -> End
+      | _ -> failwith "unreadable tag for labelled_span";
+    in
+    let offset = offset + 1 in
+    let context_size = Bytes.get_uint8 b offset in
+    let offset = offset + 1 in
+    let context = if context_size = 0 then None else Some (Bytes.sub_string b offset context_size) in
+    let offset = offset + context_size in
+    let count = BytesLabels.get_uint16_be b offset in
+    let offset = offset + 2 in
+    let line = Bytes.get_uint16_be b offset in
+    let offset = offset + 2 in
+    let filename_size = Bytes.get_uint8 b offset in
+    let offset = offset + 1 in
+    let filename = Bytes.sub_string b offset filename_size in
+    let offset = offset + filename_size in
+    assert (offset = i);
+    { kind; context; count; filename; line }
+
+  let encode b { kind; context; count; filename; line } =
+    let offset = 0 in
+    Bytes.set b offset (match kind with Begin -> 'B' | End -> 'E');
+    let offset = offset + 1 in
+    let offset = 
+      match context with
+      | None -> Bytes.set_uint8 b offset 0; offset + 1
+      | Some context ->
+          Bytes.set_uint8 b offset (String.length context);
+          let offset = offset + 1 in
+          Bytes.blit_string context 0 b offset (String.length context);
+          offset + String.length context
+    in
+    Bytes.set_uint16_be b offset count;
+    let offset = offset + 2 in
+    Bytes.set_uint16_be b offset line;
+    let offset = offset + 2 in
+    let filename_truncated_length = min (String.length filename) (Bytes.length b - offset) in
+    Bytes.set_uint8 b offset filename_truncated_length;
+    let offset = offset + 1 in
+    Bytes.blit_string filename 0 b offset filename_truncated_length;
+    let offset = offset + filename_truncated_length in
+    assert (offset <= Bytes.length b);
+    offset
+
+  let t : t Runtime_events.Type.t = Runtime_events.Type.register ~encode ~decode
+
+  type Runtime_events.User.tag += T
+  let span = Runtime_events.User.register "lwt-trace" T t
+  let emit t = Runtime_events.User.write span t
 
 end
diff --git a/src/runtime_events/lwt_runtime_events.mli b/src/runtime_events/lwt_runtime_events.mli
@@ -16,10 +16,10 @@ val unix_job_count : int Runtime_events.User.t
 val emit_job_count : int -> unit
 
 module Trace : sig
-  type labelled_span = Runtime_events.Type.span * string
-  val labelled_span : labelled_span Runtime_events.Type.t
+  type t = { kind: Runtime_events.Type.span; context: string option; count: int; filename: string; line: int; }
+  val t : t Runtime_events.Type.t
 
-  type Runtime_events.User.tag += LabelledSpan
-  val span : labelled_span Runtime_events.User.t
-  val emit_span : labelled_span -> unit
+  type Runtime_events.User.tag += T
+  val span : t Runtime_events.User.t
+  val emit : t -> unit
 end
diff --git a/test/ppx/main.ml b/test/ppx/main.ml
@@ -10,6 +10,8 @@ open Lwt
 let%lwt structure_let_result = Lwt.return_true
 [@@@ocaml.warning "+22"]
 
+let __trace_ctxt = "test" (* TODO: figure out how to make this implicit *)
+
 let suite = suite "ppx" [
   test "let"
     (fun () ->
diff --git a/test/tracing/tailgate.ml b/test/tracing/tailgate.ml
@@ -19,9 +19,27 @@ let cursor =
 
 let (_, _) = Unix.waitpid [WUNTRACED] pid
 
+let simplify_fname fname =
+  String.split_on_char '/' fname
+  |> List.rev
+  |> function
+    | [] -> assert false
+    | hd :: tl ->
+        (hd :: "/" :: List.map (fun s -> String.make 1 s.[0]) tl)
+    |> List.rev
+    |> String.concat ""
+
+let name_of { Lwt_runtime_events.Trace.context; count; filename; line; kind=_ } =
+  let base =
+    match context with
+    | None -> simplify_fname filename ^ ":" ^ string_of_int line
+    | Some c -> c
+  in
+  base ^ string_of_int count
+
 let () = Printf.printf "crash! writing trace file %s/%d.tail\n" tmpdir pid
 let () =
-  let buf_pool = Trace_fuchsia.Buf_pool.create () in
+  let buf_pool = Trace_fuchsia.Buf_pool.create ~max_size:1024 () in
   let buf = Trace_fuchsia.Buf_chain.create ~sharded:false ~buf_pool () in
   let oc = open_out_bin (Printf.sprintf "%s/%d.tail" tmpdir pid) in
   let { Trace_fuchsia.Exporter.write_bufs; flush; close } as exporter = Trace_fuchsia.Exporter.of_out_channel ~close_channel:true oc in
@@ -30,25 +48,36 @@ let () =
   let cb =
     Runtime_events.Callbacks.create ()
     |> Runtime_events.Callbacks.add_user_event
-        Lwt_runtime_events.Trace.labelled_span
+        Lwt_runtime_events.Trace.t
         (fun _ t u x ->
           match Runtime_events.User.tag u with
-            | Lwt_runtime_events.Trace.LabelledSpan -> begin
+            | Lwt_runtime_events.Trace.T -> begin
                 match x with
-                | Begin, s ->
-                    Trace_fuchsia.Writer.Event.Duration_begin.encode buf ~name:s
+                | { kind = Begin; context=_; count=_; filename; line } ->
+                    Trace_fuchsia.Writer.Event.Duration_begin.encode buf
+                    ~name:(name_of x)
                     ~t_ref:(Ref 1)
                     ~time_ns:(Runtime_events.Timestamp.to_int64 t)
-                    ~args:[]
+                    ~args:["file", A_string filename; "line", A_int line]
                     ()
-                | End, s ->
-                    Trace_fuchsia.Writer.Event.Duration_end.encode buf ~name:s
+                | { kind = End; context=_; count=_; filename; line } ->
+                    Trace_fuchsia.Writer.Event.Duration_end.encode buf
+                    ~name:(name_of x)
                     ~t_ref:(Ref 1)
                     ~time_ns:(Runtime_events.Timestamp.to_int64 t)
-                    ~args:[]
+                    ~args:["file", A_string filename; "line", A_int line]
                     ()
             end
           | _ -> ())
+    |> Runtime_events.Callbacks.add_user_event
+        Runtime_events.Type.unit
+        (fun _ t e () ->
+          match Runtime_events.User.tag e with
+          | Lwt_runtime_events.Scheduler_lap ->
+              Trace_fuchsia.Writer.Event.Instant.encode buf ~name:"lap"
+              ~t_ref:(Ref 1) ~time_ns:(Runtime_events.Timestamp.to_int64 t)
+              ~args:[] ()
+          | _ -> ())
     |> Runtime_events.Callbacks.add_user_event
         Runtime_events.Type.int
         (fun _ t u x ->
diff --git a/test/tracing/work_and_crash.ml b/test/tracing/work_and_crash.ml
@@ -2,25 +2,52 @@ let () = Runtime_events.start ()
 
 let rec ping () =
   let p = Lwt.pause () in
-  let%lwt _ping = Lwt_unix.sleep 0.01 in
+  let%lwt _ping = Lwt_unix.sleep 0.01
+  and _ping2 = Lwt_unix.sleep 0.01 in
   pong p
 and pong p =
   let%lwt _ = p in
   let%lwt _pong = Lwt_unix.sleep 0.01 in
   ping ()
 ;;
 
+let fibobo () =
+  let rec fibobo n =
+    if n <= 0 then Lwt.return 1 else
+    let%lwt left =
+      Lwt.with_value Lwt.tracing_context (Some "left") (fun () ->
+        Lwt.bind (Lwt_unix.sleep 0.001) (fun () -> Lwt.bind (Lwt.pause ()) (fun () -> fibobo (n - 1))))
+    and right =
+      Lwt.with_value Lwt.tracing_context (Some "right") (fun () ->
+        Lwt.bind (Lwt_unix.sleep 0.002) (fun () -> Lwt.bind (Lwt.pause ()) (fun () -> fibobo (n - 2))))
+    in
+    Lwt.return (left + right)
+  in
+  let%lwt () = Lwt_unix.sleep 0.02 in
+  let%lwt _ = Lwt_list.map_s fibobo (List.init 6 (fun x -> x+2)) in
+  fst (Lwt.task ()) (* never resolve *)
+;;
+
 let rec eventually_crash n =
   if n < 0 then raise Exit else
     let%lwt do_some_work = Lwt_unix.sleep 0.004 in
     ignore do_some_work;
     eventually_crash (n - 1)
 ;;
 
+let eventually_crash n =
+  Lwt.with_value Lwt.tracing_context (Some "crrrrash") (fun () -> eventually_crash n)
+
 let protek f =
-  try%lwt f () with Exit -> exit 1
+  try%lwt
+    Lwt.with_value Lwt.tracing_context (Some "protekted") f
+  with Exit -> exit 1
 
 let () = Lwt_main.run begin
   let%lwt () = Lwt.pause () in
-  Lwt.pick [ping (); protek (fun () -> eventually_crash 20)]
+  Lwt.pick [
+    (Lwt.with_value Lwt.tracing_context (Some "pingpong") ping);
+    (Lwt.with_value Lwt.tracing_context (Some "fib") fibobo);
+    protek (fun () -> eventually_crash 20)
+  ]
 end
