From 2b3f693520612a4580500f8c0b4b46f36fd4a4e1 Mon Sep 17 00:00:00 2001 From: Simon Cruanes Date: Wed, 19 Mar 2025 21:02:54 -0400 Subject: [PATCH] wip: use a concurrent hashmap for otel-trace --- .github/workflows/main.yml | 3 +- dune-project | 3 +- src/trace/ctbl_.basic.ml | 28 ++++++++++++++++++ src/trace/ctbl_.mli | 11 +++++++ src/trace/ctbl_.picos.ml | 25 ++++++++++++++++ src/trace/dune | 10 ++++++- src/trace/opentelemetry_trace.ml | 51 ++++++++++++++++++++++++++++---- 7 files changed, 122 insertions(+), 9 deletions(-) create mode 100644 src/trace/ctbl_.basic.ml create mode 100644 src/trace/ctbl_.mli create mode 100644 src/trace/ctbl_.picos.ml diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 35780eb4..7d495891 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -49,7 +49,8 @@ jobs: - run: opam exec -- dune build @install -p opentelemetry,opentelemetry-lwt,opentelemetry-client-ocurl,opentelemetry-cohttp-lwt,opentelemetry-client-cohttp-lwt - - run: opam install trace.0.9 + # with depopts + - run: opam install trace.0.9 picos_aux - run: opam exec -- dune build @install -p opentelemetry - run: opam install ocaml-protoc diff --git a/dune-project b/dune-project index 7f428494..b0320cb3 100644 --- a/dune-project +++ b/dune-project @@ -45,7 +45,7 @@ :with-dev-setup (>= 0.24) (< 0.25)))) - (depopts trace lwt eio) + (depopts trace lwt eio picos_aux) (conflicts (trace (< 0.9))) @@ -100,6 +100,7 @@ (odoc :with-doc) (lwt (>= "5.3")) + (cohttp (and (>= 4.0) (< 6.0))) (cohttp-lwt (and (>= "4.0.0") diff --git a/src/trace/ctbl_.basic.ml b/src/trace/ctbl_.basic.ml new file mode 100644 index 00000000..2b5d4cd2 --- /dev/null +++ b/src/trace/ctbl_.basic.ml @@ -0,0 +1,28 @@ +module Tbl = Hashtbl.Make (struct + include Int64 + + let hash = Hashtbl.hash +end) + +type 'a t = { + m: Mutex.t; + tbl: 'a Tbl.t; +} + +let create () : _ t = { m = Mutex.create (); tbl = Tbl.create 16 } + +let find (self : _ t) k : _ option = + Mutex.lock self.m; + let r = Tbl.find_opt self.tbl k in + Mutex.unlock self.m; + r + +let add (self : _ t) k v : unit = + Mutex.lock self.m; + Tbl.replace self.tbl k v; + Mutex.unlock self.m + +let remove (self : _ t) k : unit = + Mutex.lock self.m; + Tbl.remove self.tbl k; + Mutex.unlock self.m diff --git a/src/trace/ctbl_.mli b/src/trace/ctbl_.mli new file mode 100644 index 00000000..e27eec9c --- /dev/null +++ b/src/trace/ctbl_.mli @@ -0,0 +1,11 @@ +(** Concurrent hashmap *) + +type 'a t + +val create : unit -> 'a t + +val find : 'a t -> int64 -> 'a option + +val add : 'a t -> int64 -> 'a -> unit + +val remove : 'a t -> int64 -> unit diff --git a/src/trace/ctbl_.picos.ml b/src/trace/ctbl_.picos.ml new file mode 100644 index 00000000..c2b4c079 --- /dev/null +++ b/src/trace/ctbl_.picos.ml @@ -0,0 +1,25 @@ +module P_tbl = Picos_aux_htbl + +type 'a t = (int64, 'a) P_tbl.t + +let create () : _ t = + P_tbl.create + ~hashed_type: + (module struct + include Int64 + + let hash = Hashtbl.hash + end) + () + +let find (self : 'a t) (k : int64) : 'a option = + try Some (P_tbl.find_exn self k) with Not_found -> None + +let add (self : _ t) k v : unit = + let _ok = P_tbl.try_add self k v in + ignore _ok; + () + +let remove (self : _ t) k : unit = + ignore (P_tbl.try_remove self k : bool); + () diff --git a/src/trace/dune b/src/trace/dune index 9606e2a5..1055d556 100644 --- a/src/trace/dune +++ b/src/trace/dune @@ -3,4 +3,12 @@ (public_name opentelemetry.trace) (synopsis "Use opentelemetry as a collector for trace") (optional) - (libraries opentelemetry.ambient-context trace.core opentelemetry)) + (libraries + opentelemetry.ambient-context + trace.core + opentelemetry + (select + ctbl_.ml + from + (picos_aux.htbl -> ctbl_.picos.ml) + (-> ctbl_.basic.ml)))) diff --git a/src/trace/opentelemetry_trace.ml b/src/trace/opentelemetry_trace.ml index bf9a091d..be55694b 100644 --- a/src/trace/opentelemetry_trace.ml +++ b/src/trace/opentelemetry_trace.ml @@ -1,6 +1,7 @@ module Otel = Opentelemetry module Otrace = Trace_core (* ocaml-trace *) module TLS = Thread_local_storage +module Ambient_context = Opentelemetry_ambient_context open struct let spf = Printf.sprintf @@ -58,8 +59,8 @@ type Otrace.extension_event += | Ev_set_span_kind of Otrace.explicit_span * Otel.Span_kind.t | Ev_record_exn of Otrace.explicit_span * exn * Printexc.raw_backtrace -module Internal = struct - type span_begin = { +module Span_begin = struct + type t = { start_time: int64; name: string; __FILE__: string; @@ -68,13 +69,51 @@ module Internal = struct scope: Otel.Scope.t; parent: Otel.Span_ctx.t option; } +end - module Active_span_tbl = Hashtbl.Make (struct - include Int64 +module Local_info = struct + let key : Otel.Span_ctx.t list Ambient_context.key = + Ambient_context.create_key () - let hash : t -> int = Hashtbl.hash - end) + let[@inline] get_parent () : Otel.Span_ctx.t option = + match Ambient_context.get key with + | Some (p :: _) -> Some p + | _ -> None + | exception _ _ -> None + let push_parent p : unit = + match Ambient_context.get key with + | Some l -> Ambient_context. + match FLS.get_exn key with + | l -> FLS.set key (p :: l) + | exception Failure _ -> () (* not in fiber *) + | exception FLS.Not_set -> (try FLS.set key [ p ] with _ -> ()) + + let pop_parent () : unit = + match FLS.get_exn key with + | _ :: tl -> FLS.set key tl + | [] -> () + | exception Failure _ -> () (* not in fiber *) + | exception FLS.Not_set -> () +end + +module Spans = struct + module P_tbl = Ctbl_ + + let tbl : span_begin P_tbl.t = P_tbl.create () + + let find (span : Otrace.span) : span_begin option = P_tbl.find tbl span + + let enter (span : Trace.span) (sb : Span_begin.t) : unit = + P_tbl.add tbl span sb; + Local_info.push_parent (Span_begin.to_span_ctx sb) + + let exit (span : Trace.span) : unit = + ignore (P_tbl.try_remove tbl span : bool); + Local_info.pop_parent () +end + +module Internal = struct (** key to access a OTEL scope from an explicit span *) let k_explicit_scope : Otel.Scope.t Otrace.Meta_map.key = Otrace.Meta_map.Key.create ()