refactor collector API as well.

After long discussions with @ELLIOTTCABLE, the collector should only
provide a `with_`-like function for regular spans. This allows
each concurrency library to use its native mechanism for local storage,
which is almost always based on such a scope function.

src/core/collector.ml

@@ -17,18 +17,19 @@ let dummy_explicit_span : explicit_span =
     This is only relevant to implementors of tracing backends; to instrument
     your code you only need to look at the {!Trace} module. *)
 module type S = sig
-  val enter_span :
-    ?__FUNCTION__:string ->
+  val with_span :
+    __FUNCTION__:string option ->
     __FILE__:string ->
     __LINE__:int ->
     data:(string * user_data) list ->
     string ->
-    span
-  (** Enter a new span. *)
+    (span -> 'a) ->
+    'a
+  (** Run the function in a new span.
 
-  val exit_span : span -> unit
-  (** Exit given span. It can't be exited again. Spans must follow
-      a strict stack discipline on each thread. *)
+      This replaces the previous [enter_span] and [exit_span] which were too flexible
+      to be efficient to implement in async contexts.
+     @since NEXT_RELEASE *)
 
   val enter_explicit_span :
     surrounding:explicit_span option ->

src/core/trace_core.ml

@@ -13,34 +13,10 @@ let[@inline] enabled () =
   | None -> false
   | Some _ -> true
 
-let enter_span_collector_ (module C : Collector.S) ?__FUNCTION__ ~__FILE__
-    ~__LINE__ ?(data = fun () -> []) name : span =
-  let data = data () in
-  C.enter_span ?__FUNCTION__ ~__FILE__ ~__LINE__ ~data name
-
-let[@inline] enter_span ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data name : span =
-  match A.get collector with
-  | None -> Collector.dummy_span
-  | Some coll ->
-    enter_span_collector_ coll ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data name
-
-let[@inline] exit_span span : unit =
-  match A.get collector with
-  | None -> ()
-  | Some (module C) -> C.exit_span span
-
 let with_span_collector_ (module C : Collector.S) ?__FUNCTION__ ~__FILE__
     ~__LINE__ ?(data = fun () -> []) name f =
   let data = data () in
-  let sp = C.enter_span ?__FUNCTION__ ~__FILE__ ~__LINE__ ~data name in
-  match f sp with
-  | x ->
-    C.exit_span sp;
-    x
-  | exception exn ->
-    let bt = Printexc.get_raw_backtrace () in
-    C.exit_span sp;
-    Printexc.raise_with_backtrace exn bt
+  C.with_span ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name f
 
 let[@inline] with_span ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data name f =
   match A.get collector with

src/core/trace_core.mli

@@ -12,29 +12,6 @@ val enabled : unit -> bool
     This is fast, so that the traced program can check it before creating
     any span or message *)
 
-val enter_span :
-  ?__FUNCTION__:string ->
-  __FILE__:string ->
-  __LINE__:int ->
-  ?data:(unit -> (string * user_data) list) ->
-  string ->
-  span
-(** Enter a span. A span is a delimited period of time with
-    a start instant ("enter") and stop instant ("exit"), associated
-    with some metadata about the code entering/exiting some piece of code.
-
-    In particular the entrypoint comes with the location of the code
-    (you can use [__FILE__] and [__LINE__] directly in OCaml),
-    a mandatory name, and some optional metadata in a JSON-like representation.
-*)
-
-val exit_span : span -> unit
-(** Exit the span.
-
-    This should be called exactly once per span (even in case of exception).
-    Once this is called, a timestamp might be recorded,
-  and the span is considered finished and can't be used anymore *)
-
 val with_span :
   ?__FUNCTION__:string ->
   __FILE__:string ->

src/tef/trace_tef.ml

@@ -328,17 +328,19 @@ let collector ~out () : collector =
       else
         Thread.id (Thread.self ())
 
-    let enter_span ?__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_ ~data name : span =
+    let with_span ~__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_ ~data name f =
       let span = Int64.of_int (A.fetch_and_add span_id_gen_ 1) in
       let tid = get_tid_ () in
       let time_us = now_us () in
       B_queue.push events
         (E_define_span { tid; name; time_us; id = span; data });
-      span
 
-    let exit_span span : unit =
-      let time_us = now_us () in
-      B_queue.push events (E_exit_span { id = span; time_us })
+      let finally () =
+        let time_us = now_us () in
+        B_queue.push events (E_exit_span { id = span; time_us })
+      in
+
+      Fun.protect ~finally (fun () -> f span)
 
     let enter_explicit_span ~(surrounding : explicit_span option)
         ?__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_ ~data name : explicit_span =