From a5eef687c8a15f22ae8e86759ef5dbd574676ce1 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Sat, 3 Feb 2024 17:21:40 -0500
Subject: [PATCH] expose Tracing

---
 src/core/moonpool.ml          |  1 +
 src/core/moonpool.mli         |  2 ++
 src/private/tracing_.dummy.ml |  1 +
 src/private/tracing_.mli      |  1 +
 src/private/tracing_.real.ml  | 11 +++++++++++
 5 files changed, 16 insertions(+)

diff --git a/src/core/moonpool.ml b/src/core/moonpool.ml
index c69b5581..aed377ea 100644
--- a/src/core/moonpool.ml
+++ b/src/core/moonpool.ml
@@ -34,6 +34,7 @@ module Private = struct
   module Ws_deque_ = Ws_deque_
   module Suspend_ = Suspend_
   module Domain_ = Domain_
+  module Tracing_ = Tracing_
 
   let num_domains = Domain_pool_.n_domains
 end
diff --git a/src/core/moonpool.mli b/src/core/moonpool.mli
index c8049503..cb1c2991 100644
--- a/src/core/moonpool.mli
+++ b/src/core/moonpool.mli
@@ -222,6 +222,8 @@ module Private : sig
   module Domain_ = Domain_
   (** Utils for domains *)
 
+  module Tracing_ = Tracing_
+
   val num_domains : unit -> int
   (** Number of domains in the backing domain pool *)
 end
diff --git a/src/private/tracing_.dummy.ml b/src/private/tracing_.dummy.ml
index d685a5b3..ba6d0aff 100644
--- a/src/private/tracing_.dummy.ml
+++ b/src/private/tracing_.dummy.ml
@@ -3,3 +3,4 @@ let dummy_span = 0L
 let enter_span _name = dummy_span
 let exit_span = ignore
 let set_thread_name = ignore
+let with_span _ f = f dummy_span
diff --git a/src/private/tracing_.mli b/src/private/tracing_.mli
index 35379332..d4634697 100644
--- a/src/private/tracing_.mli
+++ b/src/private/tracing_.mli
@@ -1,5 +1,6 @@
 val dummy_span : int64
 val enter_span : string -> int64
 val exit_span : int64 -> unit
+val with_span : string -> (int64 -> 'a) -> 'a
 val enabled : unit -> bool
 val set_thread_name : string -> unit
diff --git a/src/private/tracing_.real.ml b/src/private/tracing_.real.ml
index f71ec418..4a928e27 100644
--- a/src/private/tracing_.real.ml
+++ b/src/private/tracing_.real.ml
@@ -12,3 +12,14 @@ let[@inline] enter_span name : int64 =
     Trace.enter_span ~__FILE__:dummy_file_ ~__LINE__:0 name
 
 let[@inline] exit_span sp = if sp <> dummy_span then Trace.exit_span sp
+
+let with_span name f =
+  let sp = enter_span name in
+  try
+    let x = f sp in
+    exit_span sp;
+    x
+  with exn ->
+    let bt = Printexc.get_raw_backtrace () in
+    exit_span sp;
+    Printexc.raise_with_backtrace exn bt