Source file `trace_runtime_events.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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
(** Simple backend that emits trace events via Runtime_events.

    This backend allows trace spans, messages, and metrics to be collected by
    external tools using the OCaml Runtime_events system. *)

open Trace_core

(* Register custom event types for strings *)
module String_type = struct
  let max_len = 1024

  let encode buf s =
    let len = min (String.length s) (max_len - 1) in
    Bytes.blit_string s 0 buf 0 len;
    len

  let decode buf len = Bytes.sub_string buf 0 len
  let ty = Runtime_events.Type.register ~encode ~decode
end

module String_int = struct
  let max_len = 1024

  let encode buf (s, i) =
    let len = min (String.length s) (max_len - 9) in
    Bytes.set_int64_le buf 0 (Int64.of_int i);
    Bytes.blit_string s 0 buf 8 len;
    len + 8

  let decode buf len =
    let i = Bytes.get_int64_le buf 0 in
    Bytes.sub_string buf 8 (len - 8), Int64.to_int i

  let ty = Runtime_events.Type.register ~encode ~decode
end

module String_float = struct
  let max_len = 1024

  let encode buf (s, f) =
    let len = min (String.length s) (max_len - 9) in
    Bytes.set_int64_le buf 0 (Int64.bits_of_float f);
    Bytes.blit_string s 0 buf 8 len;
    len + 8

  let decode buf len =
    let i = Bytes.get_int64_le buf 0 in
    Bytes.sub_string buf 8 (len - 8), Int64.float_of_bits i

  let ty = Runtime_events.Type.register ~encode ~decode
end

module Events = struct
  (* Define event tags *)
  type Runtime_events.User.tag +=
    | Tag_span_enter
    | Tag_span_exit
    | Tag_message
    | Tag_metric_int
    | Tag_metric_float

  (* Register user events *)
  let span_enter_event =
    Runtime_events.User.register "trace.span.enter" Tag_span_enter
      String_type.ty

  let span_exit_event =
    Runtime_events.User.register "trace.span.exit" Tag_span_exit String_type.ty

  let message_event =
    Runtime_events.User.register "trace.message" Tag_message String_type.ty

  let metric_int_event =
    Runtime_events.User.register "trace.metric.int" Tag_metric_int String_int.ty

  let metric_float_event =
    Runtime_events.User.register "trace.metric.float" Tag_metric_float
      String_float.ty
end

(* Span representation *)
type span_info = { name: string }
type Trace_core.span += Span_runtime_events of span_info

(* Collector state *)
type st = {
  active: bool Trace_core.Internal_.Atomic_.t;
  start_events: bool;
}

let create ?(start_events = true) () : st =
  { active = Trace_core.Internal_.Atomic_.make true; start_events }

(* Collector callbacks *)
let init (self : st) = if self.start_events then Runtime_events.start ()

let shutdown (self : st) =
  Trace_core.Internal_.Atomic_.set self.active false;
  Runtime_events.pause ()

let enabled _ _ = true

let enter_span (_self : st) ~__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_ ~level:_
    ~params:_ ~data:_ ~parent:_ name : span =
  Runtime_events.User.write Events.span_enter_event name;
  Span_runtime_events { name }

let exit_span (_self : st) sp =
  match sp with
  | Span_runtime_events info ->
    Runtime_events.User.write Events.span_exit_event info.name
  | _ -> ()

let add_data_to_span _st _sp _data =
  (* Runtime_events doesn't support adding data to spans after creation,
     so we just ignore this *)
  ()

let message (_self : st) ~level:_ ~params:_ ~data:_ ~span:_ msg : unit =
  Runtime_events.User.write Events.message_event msg

let metric (_self : st) ~level:_ ~params:_ ~data:_ name m : unit =
  match m with
  | Core_ext.Metric_int n ->
    Runtime_events.User.write Events.metric_int_event (name, n)
  | Core_ext.Metric_float f ->
    Runtime_events.User.write Events.metric_float_event (name, f)
  | _ -> ()

let extension _self ~level:_ _ev =
  (* Extension events like set_thread_name, set_process_name could be
     emitted as custom events if needed *)
  ()

(* Create collector *)
let callbacks : st Collector.Callbacks.t =
  Collector.Callbacks.make ~init ~shutdown ~enabled ~enter_span ~exit_span
    ~add_data_to_span ~message ~metric ~extension ()

let collector ?(start_events = true) () : Collector.t =
  let st = create ~start_events () in
  Collector.C_some (st, callbacks)

(* Setup function *)
let setup ?(start_events = true) () =
  Trace_core.setup_collector (collector ~start_events ())

(* Convenience wrapper *)
let with_setup ?start_events f =
  setup ?start_events ();
  Fun.protect ~finally:Trace_core.shutdown f