ppx_trace

README.md

@@ -12,6 +12,7 @@ a library or application, either by hand or via a ppx.
 - [x] messages
 - [x] counters
 - [ ] other metrics?
+- [x] ppx to help instrumentation
 
 ### Usage
 
@@ -74,6 +75,54 @@ Opening it in https://ui.perfetto.dev we get something like this:
 
 ![screenshot of perfetto UI](media/ui.png)
 
+## ppx_trace
+
+On OCaml >= 4.12, and with `ppxlib` installed, you can install `ppx_trace`.
+This is a preprocessor that will rewrite like so:
+
+```ocaml
+let%trace f x y z =
+  do_sth x;
+  do_sth y;
+  begin
+    let%trace () = "sub-span" in
+    do_sth z
+  end
+```
+
+This more or less corresponds to:
+
+```ocaml
+let f x y z =
+  let _trace_span = Trace_core.enter_span ~__FILE__ ~__LINE__ "Foo.f" in
+  match
+    do_sth x;
+    do_sth y;
+    begin
+      let _trace_span = Trace_core.enter_span ~__FILE__ ~__LINE__ "sub-span" in
+      match do_sth z with
+      | res ->
+        Trace_core.exit_span _trace_span;
+        res
+      | exception e ->
+        Trace_core.exit_span _trace_span
+        raise e
+    end;
+  with
+  | res ->
+    Trace_core.exit_span _trace_span
+    res
+  | exception e ->
+    Trace_core.exit_span _trace_span
+    raise e
+```
+
+### Dune configuration
+
+In your `library` or `executable` stanza, add: `(preprocess (pps ppx_trace))`.
+The dependency on `trace.core` is automatically added. You still need to
+configure a backend to actually do collection.
+
 ### Backends
 
 Concrete tracing or observability formats such as:

dune-project

@@ -21,6 +21,18 @@
  (tags
   (trace tracing observability profiling)))
 
+(package
+ (name ppx_trace)
+ (synopsis "ppx-based instrumentation for trace")
+ (depends
+   (ocaml (>= 4.12)) ; we use __FUNCTION__
+   ppxlib
+   (trace (= :version))
+   (trace-tef (and (= :version) :with-test))
+   dune)
+ (tags
+  (trace tracing observability profiling ppx)))
+
 (package
  (name trace-tef)
  (synopsis "A simple backend for trace, emitting Catapult/TEF JSON into a file")

ppx_trace.opam

@@ -0,0 +1,35 @@
+# This file is generated by dune, edit dune-project instead
+opam-version: "2.0"
+version: "0.5"
+synopsis: "ppx-based instrumentation for trace"
+maintainer: ["Simon Cruanes"]
+authors: ["Simon Cruanes"]
+license: "MIT"
+tags: ["trace" "tracing" "observability" "profiling" "ppx"]
+homepage: "https://github.com/c-cube/ocaml-trace"
+bug-reports: "https://github.com/c-cube/ocaml-trace/issues"
+depends: [
+  "ocaml" {>= "4.12"}
+  "ppxlib"
+  "trace" {= version}
+  "trace-tef" {= version & with-test}
+  "dune" {>= "2.9"}
+  "odoc" {with-doc}
+]
+build: [
+  ["dune" "subst"] {dev}
+  [
+    "dune"
+    "build"
+    "-p"
+    name
+    "-j"
+    jobs
+    "--promote-install-files=false"
+    "@install"
+    "@runtest" {with-test}
+    "@doc" {with-doc}
+  ]
+  ["dune" "install" "-p" name "--create-install-files" name]
+]
+dev-repo: "git+https://github.com/c-cube/ocaml-trace.git"

src/core/collector.ml

@@ -31,6 +31,23 @@ module type S = sig
       to be efficient to implement in async contexts.
      @since 0.3 *)
 
+  val enter_span :
+    __FUNCTION__:string option ->
+    __FILE__:string ->
+    __LINE__:int ->
+    data:(string * user_data) list ->
+    string ->
+    span
+  (** Enter a new implicit span. For many uses cases, {!with_span} will
+      be easier to use.
+      @since NEXT_RELEASE *)
+
+  val exit_span : span -> unit
+  (** Exit span. This should be called on the same thread
+      as the corresponding {!enter_span}, and nest properly with
+      other calls to enter/exit_span and {!with_span}.
+      @since NEXT_RELEASE *)
+
   val enter_manual_span :
     parent:explicit_span option ->
     flavor:[ `Sync | `Async ] option ->

src/core/trace_core.ml

@@ -29,6 +29,19 @@ let[@inline] with_span ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data name f =
     with_span_collector_ collector ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data name
       f
 
+let[@inline] enter_span ?__FUNCTION__ ~__FILE__ ~__LINE__
+    ?(data = data_empty_build_) name : span =
+  match A.get collector with
+  | None -> Collector.dummy_span
+  | Some (module C) ->
+    let data = data () in
+    C.enter_span ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name
+
+let[@inline] exit_span sp : unit =
+  match A.get collector with
+  | None -> ()
+  | Some (module C) -> C.exit_span sp
+
 let enter_explicit_span_collector_ (module C : Collector.S) ~parent ~flavor
     ?__FUNCTION__ ~__FILE__ ~__LINE__ ?(data = data_empty_build_) name :
     explicit_span =

src/core/trace_core.mli

@@ -34,6 +34,16 @@ val with_span :
     see {!enter_manual_span}.
 *)
 
+val enter_span :
+  ?__FUNCTION__:string ->
+  __FILE__:string ->
+  __LINE__:int ->
+  ?data:(unit -> (string * user_data) list) ->
+  string ->
+  span
+
+val exit_span : span -> unit
+
 val add_data_to_span : span -> (string * user_data) list -> unit
 (** Add structured data to the given active span (see {!with_span}).
     Behavior is not specified if the span has been exited.

src/ppx/dune

@@ -0,0 +1,8 @@
+
+(library
+ (name ppx_trace)
+ (public_name ppx_trace)
+ (kind ppx_rewriter)
+ (preprocess (pps ppxlib.metaquot))
+ (ppx_runtime_libraries trace.core)
+ (libraries ppxlib))

src/ppx/ppx_trace.ml

@@ -0,0 +1,82 @@
+open Ppxlib
+
+let location_errorf ~loc fmt =
+  Format.kasprintf
+    (fun err ->
+      raise (Ocaml_common.Location.Error (Ocaml_common.Location.error ~loc err)))
+    fmt
+
+(** {2 let expression} *)
+
+let expand_let ~ctxt (name : string) body =
+  let loc = Expansion_context.Extension.extension_point_loc ctxt in
+  Ast_builder.Default.(
+    [%expr
+      let _trace_span =
+        Trace_core.enter_span ~__FILE__ ~__LINE__ [%e estring ~loc name]
+      in
+      try
+        let res = [%e body] in
+        Trace_core.exit_span _trace_span;
+        res
+      with exn ->
+        Trace_core.exit_span _trace_span;
+        raise exn])
+
+let extension_let =
+  Extension.V3.declare "trace" Extension.Context.expression
+    (let open! Ast_pattern in
+    single_expr_payload
+      (pexp_let nonrecursive
+         (value_binding
+            ~pat:(ppat_construct (lident (string "()")) none)
+            ~expr:(estring __)
+         ^:: nil)
+         __))
+    expand_let
+
+let rule_let = Ppxlib.Context_free.Rule.extension extension_let
+
+(** {2 Toplevel binding} *)
+
+let expand_top_let ~ctxt rec_flag (vbs : _ list) =
+  let loc = Expansion_context.Extension.extension_point_loc ctxt in
+  Ast_builder.Default.(
+    (* go in functions, and add tracing around the body *)
+    let rec push_into_fun (e : expression) : expression =
+      match e.pexp_desc with
+      | Pexp_fun (lbl, lbl_expr, pat, body) ->
+        pexp_fun ~loc:e.pexp_loc lbl lbl_expr pat @@ push_into_fun body
+      | _ ->
+        [%expr
+          let _trace_span =
+            Trace_core.enter_span ~__FILE__ ~__LINE__ __FUNCTION__
+          in
+          match [%e e] with
+          | res ->
+            Trace_core.exit_span _trace_span;
+            res
+          | exception exn ->
+            let bt = Printexc.get_raw_backtrace () in
+            Trace_core.exit_span _trace_span;
+            Printexc.raise_with_backtrace exn bt]
+    in
+
+    let tr_vb (vb : value_binding) : value_binding =
+      let expr = push_into_fun vb.pvb_expr in
+      { vb with pvb_expr = expr }
+    in
+
+    let vbs = List.map tr_vb vbs in
+    pstr_value ~loc rec_flag vbs)
+
+let extension_top_let =
+  Extension.V3.declare "trace" Extension.Context.structure_item
+    (let open! Ast_pattern in
+    pstr (pstr_value __ __ ^:: nil))
+    expand_top_let
+
+let rule_top_let = Ppxlib.Context_free.Rule.extension extension_top_let
+
+let () =
+  Driver.register_transformation ~rules:[ rule_let; rule_top_let ] "ppx_trace"

src/tef/trace_tef.ml

@@ -386,19 +386,34 @@ let collector ~out () : collector =
       else
         Thread.id (Thread.self ())
 
-    let with_span ~__FUNCTION__:fun_name ~__FILE__:_ ~__LINE__:_ ~data name f =
+    let[@inline] enter_span_ ~fun_name ~data name : span =
       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; fun_name; data });
+      span
 
-      let finally () =
-        let time_us = now_us () in
-        B_queue.push events (E_exit_span { id = span; time_us })
-      in
+    let enter_span ~__FUNCTION__:fun_name ~__FILE__:_ ~__LINE__:_ ~data name :
+        span =
+      enter_span_ ~fun_name ~data name
+
+    let exit_span span : unit =
+      let time_us = now_us () in
+      B_queue.push events (E_exit_span { id = span; time_us })
+
+    (* re-raise exception with its backtrace *)
+    external reraise : exn -> 'a = "%reraise"
 
-      Fun.protect ~finally (fun () -> f span)
+    let with_span ~__FUNCTION__:fun_name ~__FILE__:_ ~__LINE__:_ ~data name f =
+      let span = enter_span_ ~fun_name ~data name in
+      try
+        let x = f span in
+        exit_span span;
+        x
+      with exn ->
+        exit_span span;
+        reraise exn
 
     let add_data_to_span span data =
       if data <> [] then B_queue.push events (E_add_data { id = span; data })

test/dune

@@ -2,5 +2,12 @@
 (test
   (name t1)
   (package trace-tef)
+  (modules t1)
   (libraries trace trace-tef))
 
+(test
+  (name t2)
+  (package ppx_trace)
+  (modules t2)
+  (preprocess (pps ppx_trace))
+  (libraries trace-tef))