fix lwt test by providing a local ambient span provider

it's too backend-specific, it would clash eg with the opentelemetry
ambient span provider. Better let the user know this.

.github/workflows/format.yml

@@ -1,28 +1,17 @@
 name: format
 
 on:
-  pull_request:
   push:
     branches:
       - main
+  pull_request:
 
 jobs:
   format:
     name: format
-    strategy:
-      matrix:
-        ocaml-compiler:
-            - '5.3'
     runs-on: 'ubuntu-latest'
+    container: ghcr.io/c-cube/c-cube-commmon/ci-doc-5.3:latest
     steps:
-      - uses: actions/checkout@main
+      - uses: actions/checkout@v6
-      - name: Use OCaml ${{ matrix.ocaml-compiler }}
-        uses: ocaml/setup-ocaml@v3
-        with:
-          ocaml-compiler: ${{ matrix.ocaml-compiler }}
-          dune-cache: true
-          allow-prerelease-opam: true
-
-      - run: opam install ocamlformat.0.27.0
       - run: opam exec -- make format-check
 

.github/workflows/gh-pages.yml

@@ -11,26 +11,19 @@ jobs:
     runs-on: ubuntu-latest
     permissions:
       contents: write
+    container: ghcr.io/c-cube/c-cube-commmon/ci-doc-5.3:latest
     steps:
-      - uses: actions/checkout@main
+      - uses: actions/checkout@v6
 
-      - name: Use OCaml
+      - run: opam pin .#HEAD -y -n
-        uses: ocaml/setup-ocaml@v3
-        with:
-          ocaml-compiler: '5.1.x'
-          allow-prerelease-opam: true
-          dune-cache: true
-
-      - run: opam pin odoc 3.1.0 -y -n
-        # crash with 2.4, see https://github.com/ocaml/odoc/issues/1066
       - name: Deps
-        run: opam install odig thread-local-storage trace trace-tef trace-fuchsia ppx_trace
+        run: opam install thread-local-storage lwt trace trace-tef trace-fuchsia ppx_trace
 
       - name: Build
         run: opam exec -- odig odoc --cache-dir=_doc/ trace trace-tef trace-fuchsia ppx_trace
 
       - name: Deploy
-        uses: peaceiris/actions-gh-pages@v3
+        uses: peaceiris/actions-gh-pages@v4
         with:
           github_token: ${{ secrets.GITHUB_TOKEN }}
           publish_dir: ./_doc/html

.github/workflows/main.yml

@@ -1,4 +1,4 @@
-name: Build and Test
+name: build
 
 on:
   push:
@@ -8,28 +8,21 @@ on:
 
 jobs:
   run:
-    name: build
+    name: build  # build+test on various versions of OCaml, on linux
+    timeout-minutes: 15
     strategy:
       fail-fast: true
       matrix:
-        os:
+        container:
-            - ubuntu-latest
+          - ghcr.io/c-cube/c-cube-commmon/ci-4.08:latest
-              #- macos-latest
+          - ghcr.io/c-cube/c-cube-commmon/ci-4.14:latest
-              #- windows-latest
+          - ghcr.io/c-cube/c-cube-commmon/ci-5.0:latest
-        ocaml-compiler:
+          - ghcr.io/c-cube/c-cube-commmon/ci-5.4:latest
-            - '4.08.x'
-            - '4.14.x'
-            - '5.3.x'
 
-    runs-on: ${{ matrix.os }}
+    runs-on: ubuntu-latest
+    container: ${{ matrix.container }}
     steps:
-      - uses: actions/checkout@main
+      - uses: actions/checkout@v6
-      - name: Use OCaml ${{ matrix.ocaml-compiler }}
-        uses: ocaml/setup-ocaml@v3
-        with:
-          ocaml-compiler: ${{ matrix.ocaml-compiler }}
-          dune-cache: true
-          allow-prerelease-opam: true
 
       # check that trace compiles with no optional deps
       - run: opam install -t trace --deps-only
@@ -37,24 +30,16 @@ jobs:
 
       # install all packages
       - run: opam install -t trace trace-tef trace-fuchsia --deps-only
-      - run: opam install ppx_trace --deps-only # no tests
+      - run: opam install ppx_trace --deps-only
-        if: matrix.ocaml-compiler != '4.08.x'
+        if: ${{ !contains(matrix.container, 'ci-4.08') }}
       - run: opam exec -- dune build '@install' -p trace,trace-tef,trace-fuchsia
       - run: opam exec -- dune build '@install' -p ppx_trace
-        if: matrix.ocaml-compiler != '4.08.x'
+        if: ${{ !contains(matrix.container, 'ci-4.08') }}
       - run: opam exec -- dune runtest -p trace
       - run: opam install trace
       - run: opam exec -- dune runtest -p trace-tef,trace-fuchsia
 
       # with depopts
-      - run: opam install hmap thread-local-storage
+      - run: opam install thread-local-storage lwt mtime
       - run: opam exec -- dune build '@install' -p trace,trace-tef,trace-fuchsia
-
+      - run: opam exec -- dune runtest -p trace,trace-tef,trace-fuchsia
-      - run: opam install picos_aux
-        if: matrix.ocaml-compiler != '4.08.x' && matrix.ocaml-compiler != '4.12.x'
-      - run: opam exec -- dune build '@install' -p trace,trace-tef,trace-fuchsia
-        if: matrix.ocaml-compiler != '4.08.x' && matrix.ocaml-compiler != '4.12.x'
-
-      - run: opam install mtime
-      - run: opam exec -- dune build '@install' -p trace,trace-tef,trace-fuchsia
-

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,50 @@
+(** Optional interface with lwt.
+
+    This provides a Lwt-friendly variant of {!Trace_core.with_span}.
+
+    To track the ambient span, which can be useful for some backends or for
+    async spans, you can use [Ambient_context_lwt], a collector-specific ambient
+    span provider (eg [Opentelemetry_trace] comes with one), or define your own
+    like this:
+
+    {[
+      (* new key to track the current span in lwt context *)
+      let k_ambient_span : Trace_core.span Lwt.key = Lwt.new_key ()
+
+      (* ambient span provider, install it using Trace_core.set_ambient_context_provider] *)
+      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));
+            } )
+    ]}
+
+    @since NEXT_RELEASE *)
+
+open Trace_core
+
+(** [with_span name f] enters a span [sp], calls [f sp] which returns a Lwt
+    promise, and make sure to exit the span [sp] when [f sp] exits (or fails).
+
+    This is similar to {!Trace_core.with_span} but it respects the promise
+    semnatics of Lwt (ie. it doesn't exit the span immediately instead of
+    waiting for the promise to end).
+
+    To track the current span you still need to install a Lwt-friendly
+    {!Trace_core.Ambient_span_provider.t} (probably using {!Lwt.with_value}). *)
+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,60 @@
+(* Test ambient span tracking via Lwt key provider *)
+
+open Test_trace_common
+
+let ( let@ ) = ( @@ )
+let k_ambient_span : Trace_core.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 () =
+  print_endline "=== ambient span (Lwt) ===";
+  Trace_core.set_ambient_context_provider 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.
+  test_trace_common
-  ; the dummy case just echoes back the expected file.
   (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: """
+description:
-ocaml-trace can be used to instrument libraries and programs with low overhead.
+  "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."
-
-              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"}
 ]