trace.util: add the thread and time utils, add multi_collector, add span_id64 and trace_id64

utils for collectors in general really.

src/util/dune

@@ -1,14 +1,21 @@
 (library
- (public_name trace.private.util)
+ (public_name trace.util)
- (synopsis "internal utilities for trace. No guarantees of stability.")
+ (synopsis "Utilities for trace and for collectors")
- (name trace_private_util)
+ (name trace_util)
- (optional) ; depends on mtime
  (libraries
   trace.core
-  mtime
+  (select
-  mtime.clock.os
+   thread_util.ml
-  unix
+   from
-  threads
+   (threads -> thread_util.real.ml)
+   (-> thread_util.dummy.ml))
+  (select
+   time_util.ml
+   from
+   (mtime mtime.clock.os -> time_util.mtime.ml)
+   (mtime mtime.clock.jsoo -> time_util.mtime.ml)
+   (unix -> time_util.unix.ml)
+   (-> time_util.dummy.ml))
   (select
    domain_util.ml
    from

src/util/multi_collector.ml

@@ -0,0 +1,111 @@
+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__ ~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__ ~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 ~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 ~span ~params ~data msg
+    done
+    [@ocaml.warning "-8"]
+
+  let counter_int st ~params ~data name n =
+    for i = 0 to Array.length st - 1 do
+      let (Collector.C_some (st, cb)) = Array.get st i in
+      cb.counter_int st ~params ~data name n
+    done
+    [@ocaml.warning "-8"]
+
+  let counter_float st ~params ~data name n =
+    for i = 0 to Array.length st - 1 do
+      let (Collector.C_some (st, cb)) = Array.get st i in
+      cb.counter_float st ~params ~data name n
+    done
+    [@ocaml.warning "-8"]
+
+  let extension st 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 ev
+    done
+    [@ocaml.warning "-8"]
+
+  let combine_cb : st Collector.Callbacks.t =
+    {
+      Collector.Callbacks.init;
+      enter_span;
+      exit_span;
+      message;
+      add_data_to_span;
+      counter_int;
+      counter_float;
+      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 ]

src/util/multi_collector.mli

@@ -0,0 +1,11 @@
+(** Combine multiple collectors into one *)
+
+open Trace_core
+
+val combine_l : Collector.t list -> Collector.t
+(** Combine multiple collectors, ie return a collector that forwards to every
+    collector in the list. *)
+
+val combine : Collector.t -> Collector.t -> Collector.t
+(** [combine s1 s2] is a collector that forwards every call to [s1] and [s2]
+    both. *)

src/util/span_id64.ml

@@ -0,0 +1,24 @@
+open struct
+  module A = Trace_core.Internal_.Atomic_
+end
+
+type t = int64
+
+module Gen : sig
+  type t
+
+  val create : unit -> t
+  val gen : t -> int64
+end = struct
+  type t = int A.t
+
+  let create () = A.make 0
+  let[@inline] gen self : int64 = A.fetch_and_add self 1 |> Int64.of_int
+end
+
+module Trace_id_generator = struct
+  type t = int A.t
+
+  let create () = A.make 0
+  let[@inline] gen self = A.fetch_and_add self 1 |> Int64.of_int
+end

src/util/thread_util.dummy.ml

@@ -0,0 +1 @@
+let[@inline] get_tid () = 0

src/util/thread_util.mli

@@ -0,0 +1,2 @@
+val get_tid : unit -> int
+(** Get current thread ID *)

src/util/thread_util.real.ml

@@ -0,0 +1 @@
+let[@inline] get_tid () = Thread.id @@ Thread.self ()

src/util/time_util.dummy.ml

@@ -0,0 +1 @@
+let[@inline] get_time_ns () : int64 = Sys.time () *. 1e9

src/util/time_util.mli

@@ -0,0 +1,3 @@
+val get_time_ns : unit -> int64
+(** Get current time in nanoseconds. The beginning point is unspecified, and
+    this is assumed to be best-effort monotonic. Ideally, use [mtime]. *)

src/util/time_util.mtime.ml

@@ -0,0 +1,3 @@
+let[@inline] get_time_ns () : int64 =
+  let t = Mtime_clock.now () in
+  Mtime.to_uint64_ns t

src/util/time_util.unix.ml

@@ -0,0 +1,3 @@
+let[@inline] get_time_ns () : int64 =
+  let t = Unix.gettimeofday () in
+  Int64.of_float (t *. 1e9)

src/util/trace_id64.ml

@@ -0,0 +1,17 @@
+open struct
+  module A = Trace_core.Internal_.Atomic_
+end
+
+type t = int64
+
+module Gen : sig
+  type t
+
+  val create : unit -> t
+  val gen : t -> int64
+end = struct
+  type t = int A.t
+
+  let create () = A.make 0
+  let[@inline] gen self : int64 = A.fetch_and_add self 1 |> Int64.of_int
+end