Source file trace_tef_tldrs.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
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/tldrs.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;  (** Unix socket address *)
  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 =
      (* normalize path so the daemon knows what we're talking about *)
      let path =
        if Filename.is_relative path then
          Filename.concat (Unix.getcwd ()) path
        else
          path
      in
      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 subscriber_ (client : as_client) : Trace_subscriber.t =
  (* 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 `tldrs` 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 () =
    (try flush out with _ -> ());
    try Unix.close sock with _ -> ()
  in

  fpf out "OPEN %s\n%!" client.trace_id;
  (* ask the collector to emit the trace in a user-chosen file, perhaps *)
  Option.iter
    (fun file -> fpf out "EMIT_TEF_AT_EXIT %s\n" file)
    client.emit_tef_at_exit;

  Trace_tef.Private_.subscriber_jsonl ~finally ~out:(`Output out) ()

let subscriber ~out () =
  let role = find_role ~out () in
  match role with
  | None -> assert false
  | Some c -> subscriber_ c

let collector ~out () : collector =
  let role = find_role ~out () in
  match role with
  | None -> assert false
  | Some c -> subscriber_ c |> Trace_subscriber.collector

let setup ?(out = `Env) () =
  let role = find_role ~out () in
  match role with
  | None -> ()
  | Some c ->
    Trace_core.setup_collector @@ Trace_subscriber.collector @@ subscriber_ c

let with_setup ?out () f =
  setup ?out ();
  Fun.protect ~finally:Trace_core.shutdown f

module Private_ = struct
  include Trace_tef.Private_
end