wip: use a concurrent hashmap for otel-trace

.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

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")

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

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

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);
+  ()

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))))

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 ()