diff --git a/src/tef-tldr/dune b/src/tef-tldr/dune new file mode 100644 index 0000000..c0602cd --- /dev/null +++ b/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)) diff --git a/src/tef-tldr/trace_tef_tldr.ml b/src/tef-tldr/trace_tef_tldr.ml new file mode 100644 index 0000000..a712cf2 --- /dev/null +++ b/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 diff --git a/src/tef-tldr/trace_tef_tldr.mli b/src/tef-tldr/trace_tef_tldr.mli new file mode 100644 index 0000000..dc67648 --- /dev/null +++ b/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 + +(**/**)