add Trace_lwt optional library (#47)

* add Trace_lwt optional library

provides a with_span that's lwt friendly, as well as a ambient span
provider.

* test: extract recorder used for ambient tests

* tests for trace.lwt

* opam

dune-project

@@ -22,13 +22,14 @@
  (synopsis
   "A lightweight stub for tracing/observability, agnostic in how data is collected")
  (description
-  "ocaml-trace can be used to instrument libraries and programs with low overhead.\n\n              It doesn't do any IO unless a collector is plugged in, which only\n              the final executable should do.")
+  "ocaml-trace can be used to instrument libraries and programs with low overhead. It doesn't do any IO unless a collector is plugged in, which only the final executable should do.")
  (depends
   (ocaml
    (>= 4.08))
   dune)
  (depopts
   unix
+  lwt
   (thread-local-storage (>= 0.2))
   (mtime
    (>= 2.0)))

src/lwt/dune

@@ -0,0 +1,6 @@
+(library
+ (name trace_lwt)
+ (public_name trace.lwt)
+ (optional) ; lwt
+ (libraries trace.core lwt)
+ (synopsis "Interface with lwt"))

src/lwt/trace_lwt.ml

@@ -0,0 +1,30 @@
+(** Optional interface with lwt.
+
+    @since NEXT_RELEASE *)
+
+open Trace_core
+
+let k_ambient_span : span Lwt.key = Lwt.new_key ()
+
+let ambient_span_provider : Trace_core.Ambient_span_provider.t =
+  ASP_some
+    ( (),
+      {
+        get_current_span = (fun () -> Lwt.get k_ambient_span);
+        with_current_span_set_to =
+          (fun () span f ->
+            Lwt.with_value k_ambient_span (Some span) (fun () -> f span));
+      } )
+
+let with_span ?level ?__FUNCTION__ ~__FILE__ ~__LINE__ ?parent ?params ?data
+    name (f : span -> 'a Lwt.t) : 'a Lwt.t =
+  if Trace_core.enabled () then (
+    let span =
+      Trace_core.enter_span ?level ?__FUNCTION__ ~__FILE__ ~__LINE__ ?parent
+        ?params ?data name
+    in
+    let fut = Trace_core.with_current_span_set_to span f in
+    Lwt.on_termination fut (fun () -> Trace_core.exit_span span);
+    fut
+  ) else
+    f Trace_core.Collector.dummy_span

test/common/dune

@@ -0,0 +1,3 @@
+(library
+ (name test_trace_common)
+ (libraries trace.core))

test/common/test_trace_common.ml

@@ -0,0 +1,32 @@
+(** Shared test utilities for ambient span tests. *)
+
+type Trace_core.span += Named of string
+
+let make_recorder () =
+  let open Trace_core.Collector in
+  let cbs =
+    Callbacks.make
+      ~enter_span:(fun
+          ()
+          ~__FUNCTION__:_
+          ~__FILE__:_
+          ~__LINE__:_
+          ~level:_
+          ~params:_
+          ~data:_
+          ~parent:_
+          name
+        -> Named name)
+      ~exit_span:(fun () _sp -> ())
+      ~add_data_to_span:(fun () _sp _data -> ())
+      ~message:(fun () ~level:_ ~params:_ ~data:_ ~span:_ _msg -> ())
+      ~metric:(fun () ~level:_ ~params:_ ~data:_ _name _m -> ())
+      ()
+  in
+  C_some ((), cbs)
+
+let current_name () =
+  match Trace_core.current_span () with
+  | None -> "none"
+  | Some (Named s) -> s
+  | Some _ -> "<other>"

test/lwt/dune

@@ -0,0 +1,21 @@
+(executable
+ (name t_lwt)
+ (modules t_lwt)
+ (libraries
+  trace
+  test_trace_common
+  (select
+   t_lwt.ml
+   from
+   (trace.lwt lwt.unix -> t_lwt.real.ml)
+   (-> t_lwt.quine.ml))))
+
+(rule
+ (alias runtest)
+ (package trace)
+ (action
+  (progn
+   (with-stdout-to
+    t_lwt.output
+    (run %{exe:t_lwt.exe} %{dep:t_lwt.expected}))
+   (diff t_lwt.expected t_lwt.output))))

test/lwt/t_lwt.expected

@@ -0,0 +1,10 @@
+=== ambient span (Lwt) ===
+before any span: none
+sync outer: sync-outer
+sync inner: sync-inner
+after sync: none
+lwt outer: lwt-outer
+lwt inner: lwt-inner
+after lwt inner: lwt-outer
+after lwt outer: none
+after binds: bind-test

test/lwt/t_lwt.quine.ml

@@ -0,0 +1,10 @@
+(* When trace.lwt is unavailable, echo the expected file
+   so the diff always passes. The file is passed as argv[1] by the dune rule. *)
+let () =
+  let ic = open_in Sys.argv.(1) in
+  (try
+     while true do
+       print_char (input_char ic)
+     done
+   with End_of_file -> ());
+  close_in ic

test/lwt/t_lwt.real.ml

@@ -0,0 +1,49 @@
+(* Test ambient span tracking via Lwt key provider *)
+
+open Test_trace_common
+
+let ( let@ ) = ( @@ )
+
+let () =
+  print_endline "=== ambient span (Lwt) ===";
+  Trace_core.set_ambient_context_provider Trace_lwt.ambient_span_provider;
+  let coll = make_recorder () in
+  let@ () = Trace_core.with_setup_collector coll in
+
+  (* sync with_span works with the Lwt provider *)
+  Printf.printf "before any span: %s\n" (current_name ());
+
+  (let@ _outer = Trace_core.with_span ~__FILE__ ~__LINE__ "sync-outer" in
+   Printf.printf "sync outer: %s\n" (current_name ());
+   let@ _inner = Trace_core.with_span ~__FILE__ ~__LINE__ "sync-inner" in
+   Printf.printf "sync inner: %s\n" (current_name ()));
+
+  Printf.printf "after sync: %s\n" (current_name ());
+
+  (* Lwt-specific tests *)
+  Lwt_main.run
+    (let open Lwt.Syntax in
+     (* nested Trace_lwt.with_span *)
+     let* () =
+       Trace_lwt.with_span ~__FILE__ ~__LINE__ "lwt-outer" (fun _outer ->
+           Printf.printf "lwt outer: %s\n" (current_name ());
+           let* () =
+             Trace_lwt.with_span ~__FILE__ ~__LINE__ "lwt-inner" (fun _inner ->
+                 Printf.printf "lwt inner: %s\n" (current_name ());
+                 Lwt.return_unit)
+           in
+           Printf.printf "after lwt inner: %s\n" (current_name ());
+           Lwt.return_unit)
+     in
+     Printf.printf "after lwt outer: %s\n" (current_name ());
+
+     (* context survives Lwt.bind chain *)
+     let* () =
+       Trace_lwt.with_span ~__FILE__ ~__LINE__ "bind-test" (fun _sp ->
+           let* () = Lwt.return_unit in
+           let* () = Lwt.return_unit in
+           Printf.printf "after binds: %s\n" (current_name ());
+           Lwt.return_unit)
+     in
+
+     Lwt.return_unit)

test/tls/dune

@@ -3,8 +3,7 @@
  (modules t_tls)
  (libraries
   trace
-  ; avoid using a recent dune with `enabled-if`, by just using `select` instead.
-  ; the dummy case just echoes back the expected file.
+  test_trace_common
   (select
    t_tls.ml
    from

test/tls/t_tls.real.ml

@@ -1,39 +1,9 @@
 (* Test ambient span tracking via thread-local-storage provider *)
 
+open Test_trace_common
+
 let ( let@ ) = ( @@ )
 
-(* Unique span type so we can identify spans by name *)
-type Trace_core.span += Named of string
-
-let make_recorder () =
-  let open Trace_core.Collector in
-  let cbs =
-    Callbacks.make
-      ~enter_span:(fun
-          ()
-          ~__FUNCTION__:_
-          ~__FILE__:_
-          ~__LINE__:_
-          ~level:_
-          ~params:_
-          ~data:_
-          ~parent:_
-          name
-        -> Named name)
-      ~exit_span:(fun () _sp -> ())
-      ~add_data_to_span:(fun () _sp _data -> ())
-      ~message:(fun () ~level:_ ~params:_ ~data:_ ~span:_ _msg -> ())
-      ~metric:(fun () ~level:_ ~params:_ ~data:_ _name _m -> ())
-      ()
-  in
-  C_some ((), cbs)
-
-let current_name () =
-  match Trace_core.current_span () with
-  | None -> "none"
-  | Some (Named s) -> s
-  | Some _ -> "<other>"
-
 let () =
   print_endline "=== ambient span (TLS) ===";
   Trace_thread_local_storage.setup ();

trace.opam

@@ -3,11 +3,8 @@ opam-version: "2.0"
 version: "0.12"
 synopsis:
   "A lightweight stub for tracing/observability, agnostic in how data is collected"
-description: """
-ocaml-trace can be used to instrument libraries and programs with low overhead.
-
-              It doesn't do any IO unless a collector is plugged in, which only
-              the final executable should do."""
+description:
+  "ocaml-trace can be used to instrument libraries and programs with low overhead. It doesn't do any IO unless a collector is plugged in, which only the final executable should do."
 maintainer: ["Simon Cruanes"]
 authors: ["Simon Cruanes"]
 license: "MIT"
@@ -21,6 +18,7 @@ depends: [
 ]
 depopts: [
   "unix"
+  "lwt"
   "thread-local-storage" {>= "0.2"}
   "mtime" {>= "2.0"}
 ]