sidekick/src/util/Profile.ml

module type BACKEND = sig
  val get_ts : unit -> float

  val emit_duration_event :
    name : string ->
    start : float ->
    end_ : float ->
    unit ->
    unit

  val emit_instant_event :
    name : string ->
    ts : float ->
    unit ->
    unit

  val teardown : unit -> unit
end

type backend = (module BACKEND)

type probe =
  | No_probe
  | Probe of {
    name: string;
    start: float;
  }

let null_probe = No_probe

(* where to print events *)
let out_ : backend option ref = ref None

let[@inline] enabled () = match !out_ with Some _ -> true | None -> false

let[@inline never] begin_with_ (module B:BACKEND) name : probe =
  Probe {name; start=B.get_ts ()}

let[@inline] begin_ name : probe =
  match !out_ with
  | None -> No_probe
  | Some b -> begin_with_ b name

let[@inline] instant name =
  match !out_ with
  | None -> ()
  | Some (module B) ->
    let now = B.get_ts() in
    B.emit_instant_event ~name ~ts:now ()

(* slow path *)
let[@inline never] exit_full_ (module B : BACKEND) name start =
  let now = B.get_ts() in
  B.emit_duration_event ~name ~start ~end_:now ()

let[@inline] exit_with_ b pb =
  match pb with
  | No_probe -> ()
  | Probe {name; start} -> exit_full_ b name start

let[@inline] exit pb =
  match pb, !out_ with
  | Probe {name;start}, Some b -> exit_full_ b name start
  | _ -> ()

let[@inline] with_ name f =
  match !out_ with
  | None -> f()
  | Some b ->
    let pb = begin_with_ b name in
    try
      let x = f() in
      exit_with_ b pb;
      x
    with e ->
      exit_with_ b pb;
      raise e

let[@inline] with1 name f x =
  match !out_ with
  | None -> f x
  | Some b ->
    let pb = begin_with_ b name in
    try
      let res = f x in
      exit_with_ b pb;
      res
    with e ->
      exit_with_ b pb;
      raise e

let[@inline] with2 name f x y =
  with_ name (fun () -> f x y)

module Control = struct
  let setup b =
    assert (!out_ = None);
    out_ := b

  let teardown () =
    match !out_ with
    | None -> ()
    | Some (module B) ->
      out_ := None;
      B.teardown()
end