Source file metrics.ml

let enabled = ref false

let enable () = enabled := true

module Timer = struct
  module Measure = struct
    type t =
      { cumulative_time : float
      ; count : int
      }
  end

  type t =
    { start_time : float
    ; tag : string
    ; mutable stopped : bool
    }

  let aggregate = ref String.Map.empty

  let aggregated_timers () = !aggregate

  let update_aggregate { start_time; tag; stopped } =
    if not stopped then
      aggregate :=
        String.Map.update !aggregate tag ~f:(function
          | Some { Measure.cumulative_time; count } ->
            Some
              { Measure.cumulative_time =
                  cumulative_time +. (Unix.gettimeofday () -. start_time)
              ; count = count + 1
              }
          | None ->
            Some
              { Measure.cumulative_time = Unix.gettimeofday () -. start_time
              ; count = 1
              })

  let start tag = { start_time = Unix.gettimeofday (); tag; stopped = false }

  let stop t =
    match !enabled with
    | false ->
      Code_error.raise "Tried to stop a previously stopped timer"
        [ ("tag", String t.tag) ]
    | true ->
      update_aggregate t;
      t.stopped <- true

  let record tag ~f =
    match !enabled with
    | false -> f ()
    | true ->
      let start_time = Unix.gettimeofday () in
      Exn.protect ~f ~finally:(fun () ->
          update_aggregate { tag; start_time; stopped = false })
end

let reset () = Timer.aggregate := String.Map.empty