feat: add trace-tef.tldr for tracing multiple processes

src/tef-tldr/dune

@@ -0,0 +1,6 @@
+
+(library
+  (name trace_tef_tldr)
+  (public_name trace-tef.tldr)
+  (synopsis "Multiprocess tracing using the `tldr` daemon")
+  (libraries trace.core trace.private.util trace-tef unix threads))

src/tef-tldr/trace_tef_tldr.ml

@@ -0,0 +1,121 @@
+open Trace_core
+
+let spf = Printf.sprintf
+let fpf = Printf.fprintf
+
+type output = [ `File of string ]
+
+(** Env variable used to communicate to subprocesses, which trace ID to use *)
+let env_var_trace_id = "TRACE_TEF_TLDR_TRACE_ID"
+
+(** Env variable used to communicate to subprocesses, which trace ID to use *)
+let env_var_unix_socket = "TRACE_TEF_TLDR_SOCKET"
+
+let get_unix_socket () =
+  match Sys.getenv_opt env_var_unix_socket with
+  | Some s -> s
+  | None ->
+    let s = "/tmp/tldr.socket" in
+    (* children must agree on the socket file *)
+    Unix.putenv env_var_unix_socket s;
+    s
+
+type as_client = {
+  trace_id: string;
+  socket: string;
+  emit_tef_at_exit: string option;
+      (** For parent, ask daemon to emit traces here *)
+}
+
+type role = as_client option
+
+let to_hex (s : string) : string =
+  let open String in
+  let i_to_hex (i : int) =
+    if i < 10 then
+      Char.chr (i + Char.code '0')
+    else
+      Char.chr (i - 10 + Char.code 'a')
+  in
+
+  let res = Bytes.create (2 * length s) in
+  for i = 0 to length s - 1 do
+    let n = Char.code (get s i) in
+    Bytes.set res (2 * i) (i_to_hex ((n land 0xf0) lsr 4));
+    Bytes.set res ((2 * i) + 1) (i_to_hex (n land 0x0f))
+  done;
+  Bytes.unsafe_to_string res
+
+let create_trace_id () : string =
+  let now = Unix.gettimeofday () in
+  let rand = Random.State.make_self_init () in
+
+  let rand_bytes = Bytes.create 16 in
+  for i = 0 to Bytes.length rand_bytes - 1 do
+    Bytes.set rand_bytes i (Random.State.int rand 256 |> Char.chr)
+  done;
+  (* convert to hex *)
+  spf "tr-%d-%s" (int_of_float now) (to_hex @@ Bytes.unsafe_to_string rand_bytes)
+
+(** Find what this particular process has to do wrt tracing *)
+let find_role ~out () : role =
+  match Sys.getenv_opt env_var_trace_id with
+  | Some trace_id ->
+    Some { trace_id; emit_tef_at_exit = None; socket = get_unix_socket () }
+  | None ->
+    let write_to_file path =
+      let trace_id = create_trace_id () in
+      Unix.putenv env_var_trace_id trace_id;
+      { trace_id; emit_tef_at_exit = Some path; socket = get_unix_socket () }
+    in
+
+    (match out with
+    | `File path -> Some (write_to_file path)
+    | `Env ->
+      (match Sys.getenv_opt "TRACE" with
+      | Some ("1" | "true") -> Some (write_to_file "trace.json")
+      | Some path -> Some (write_to_file path)
+      | None -> None))
+
+let collector_ (client : as_client) : collector =
+  (* connect to unix socket *)
+  let sock = Unix.socket Unix.PF_UNIX Unix.SOCK_STREAM 0 in
+  (try Unix.connect sock (Unix.ADDR_UNIX client.socket)
+   with exn ->
+     failwith
+     @@ spf "Could not open socket to `tldr` demon at %S: %s" client.socket
+          (Printexc.to_string exn));
+  let out = Unix.out_channel_of_descr sock in
+
+  (* what to do when the collector shuts down *)
+  let finally () =
+    (* ask the collector to emit the trace in a user-chosen file, perhaps *)
+    Option.iter
+      (fun file -> fpf out "EMIT_TEF %s\n" file)
+      client.emit_tef_at_exit;
+    (try flush out with _ -> ());
+    try Unix.close sock with _ -> ()
+  in
+
+  fpf out "OPEN %s\n%!" client.trace_id;
+  Trace_tef.Internal_.collector_jsonl ~finally ~out:(`Output out) ()
+
+let collector ~out () : collector =
+  let role = find_role ~out () in
+  match role with
+  | None -> assert false
+  | Some c -> collector_ c
+
+let setup ?(out = `Env) () =
+  let role = find_role ~out () in
+  match role with
+  | None -> ()
+  | Some c -> Trace_core.setup_collector @@ collector_ c
+
+let with_setup ?out () f =
+  setup ?out ();
+  Fun.protect ~finally:Trace_core.shutdown f
+
+module Internal_ = struct
+  include Trace_tef.Internal_
+end

src/tef-tldr/trace_tef_tldr.mli

@@ -0,0 +1,40 @@
+val collector : out:[ `File of string ] -> unit -> Trace_core.collector
+(** Make a collector that writes into the given output.
+    See {!setup} for more details. *)
+
+type output = [ `File of string ]
+(** Output for tracing.
+    - [`File "foo"] will enable tracing and print events into file
+        named "foo". The file is only written at exit.
+*)
+
+val setup : ?out:[ output | `Env ] -> unit -> unit
+(** [setup ()] installs the collector depending on [out].
+
+    @param out can take different values:
+    - regular {!output} value to specify where events go
+    - [`Env] will enable tracing if the environment
+      variable "TRACE" is set.
+
+      - If it's set to "1", then the file is "trace.json".
+      - If it's set to "stdout", then logging happens on stdout (since 0.2)
+      - If it's set to "stderr", then logging happens on stdout (since 0.2)
+      - Otherwise, if it's set to a non empty string, the value is taken
+        to be the file path into which to write.
+*)
+
+val with_setup : ?out:[ output | `Env ] -> unit -> (unit -> 'a) -> 'a
+(** [with_setup () f] (optionally) sets a collector up, calls [f()],
+    and makes sure to shutdown before exiting.
+*)
+
+(**/**)
+
+module Internal_ : sig
+  val mock_all_ : unit -> unit
+  (** use fake, deterministic timestamps, TID, PID *)
+
+  val on_tracing_error : (string -> unit) ref
+end
+
+(**/**)