Source file multi_collector.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
open Trace_core

open struct
  type st = Collector.t array
  type span += Span_combine of span array

  let init st =
    for i = 0 to Array.length st - 1 do
      let (Collector.C_some (st, cb)) = Array.get st i in
      cb.init st
    done
    [@ocaml.warning "-8"]

  let shutdown st =
    for i = 0 to Array.length st - 1 do
      let (Collector.C_some (st, cb)) = Array.get st i in
      cb.shutdown st
    done
    [@ocaml.warning "-8"]

  let enter_span st ~__FUNCTION__ ~__FILE__ ~__LINE__ ~level ~params ~data
      ~parent name : span =
    let spans =
      Array.map
        (fun [@ocaml.warning "-8"] coll ->
          let (Collector.C_some (st, cb)) = coll in
          cb.enter_span st ~__FUNCTION__ ~__FILE__ ~__LINE__ ~level ~params
            ~data ~parent name)
        st
    in
    Span_combine spans

  let exit_span st span =
    match span with
    | Span_combine spans ->
      assert (Array.length spans = Array.length st);
      for i = 0 to Array.length st - 1 do
        let (Collector.C_some (st, cb)) = Array.get st i in
        cb.exit_span st spans.(i)
      done
      [@ocaml.warning "-8"]
    | _ -> ()

  let add_data_to_span st span data =
    match span with
    | Span_combine spans when data <> [] ->
      assert (Array.length spans = Array.length st);
      for i = 0 to Array.length st - 1 do
        let (Collector.C_some (st, cb)) = Array.get st i in
        cb.add_data_to_span st spans.(i) data
      done
      [@ocaml.warning "-8"]
    | _ -> ()

  let message st ~level ~params ~data ~span msg =
    for i = 0 to Array.length st - 1 do
      let (Collector.C_some (st, cb)) = Array.get st i in
      cb.message st ~level ~span ~params ~data msg
    done
    [@ocaml.warning "-8"]

  let metric st ~level ~params ~data name m =
    for i = 0 to Array.length st - 1 do
      let (Collector.C_some (st, cb)) = Array.get st i in
      cb.metric st ~level ~params ~data name m
    done
    [@ocaml.warning "-8"]

  let extension st ~level ev : unit =
    for i = 0 to Array.length st - 1 do
      let (Collector.C_some (st, cb)) = Array.get st i in
      cb.extension st ~level ev
    done
    [@ocaml.warning "-8"]

  let enabled st level : bool =
    Array.exists
      (fun (Collector.C_some (st, cb)) -> cb.enabled st level)
      st [@ocaml.warning "-8"]

  let combine_cb : st Collector.Callbacks.t =
    {
      Collector.Callbacks.init;
      enter_span;
      exit_span;
      enabled;
      message;
      add_data_to_span;
      metric;
      extension;
      shutdown;
    }
end

let combine_l (cs : Collector.t list) : Collector.t =
  let cs =
    List.filter
      (function
        | Collector.C_none -> false
        | Collector.C_some _ -> true)
      cs
  in
  match cs with
  | [] -> C_none
  | [ c ] -> c
  | _ -> C_some (Array.of_list cs, combine_cb)

let combine (s1 : Collector.t) (s2 : Collector.t) : Collector.t =
  combine_l [ s1; s2 ]