feature: allocation profiling

opam/dune.opam

@@ -57,6 +57,7 @@ depends: [
   "ocamlfind" { with-dev-setup & os != "win32" }
   "odoc" { with-dev-setup & >= "3.2.0" & os != "win32" }
   "ppx_expect" { with-dev-setup & >= "v0.17" & os != "win32" }
+  "ppx_optcomp" {with-dev-setup}
   "spawn" { with-dev-setup }
   "ppx_inline_test" { with-dev-setup & os != "win32" }
   "ppxlib" { with-dev-setup & >= "0.37.0" & os != "win32" }

opam/dune.opam.template

@@ -18,6 +18,7 @@ depends: [
   "ocamlfind" { with-dev-setup & os != "win32" }
   "odoc" { with-dev-setup & >= "3.2.0" & os != "win32" }
   "ppx_expect" { with-dev-setup & >= "v0.17" & os != "win32" }
+  "ppx_optcomp" {with-dev-setup}
   "spawn" { with-dev-setup }
   "ppx_inline_test" { with-dev-setup & os != "win32" }
   "ppxlib" { with-dev-setup & >= "0.37.0" & os != "win32" }

src/dune_engine/build_system.ml

@@ -1121,6 +1121,7 @@ let run f =
         ~restart:(Option.is_some files)
         ~files
         ~start);
+    Dune_trace.reset_alloc_profile ();
     (* CR-someday amokhov: Currently we invalidate cached timestamps on every
        incremental rebuild. This conservative approach helps us to work around
        some [mtime] resolution problems (e.g. on Mac OS). It would be nice to
@@ -1165,6 +1166,9 @@ let run f =
     (match Scheduler.finish_build ~stop with
      | Restarting -> ()
      | Finished { restart_duration } ->
+       let alloc_summary =
+         Dune_trace.capture_alloc_profile (`Build (Run_id.to_int run_id))
+       in
        Dune_trace.emit Build (fun () ->
          Dune_trace.Event.watch_build_finish
            ~run_id:(Run_id.to_int run_id)
@@ -1174,7 +1178,8 @@ let run f =
               | Error `Already_reported -> `Failure)
            ~start
            ~stop
-           ~restart_duration));
+           ~restart_duration);
+       Option.iter alloc_summary ~f:Dune_trace.always_emit);
     outcome
   in
   Fiber.Mutex.with_lock State.build_mutex ~f

src/dune_trace/alloc.ml

@@ -0,0 +1,291 @@
+open Stdune
+
+module Memprof = struct
+  [@@@ocaml.warning "-69-32-34"]
+
+  type t = unit
+  type allocation_source = |
+
+  type allocation =
+    { n_samples : int
+    ; size : int
+    ; source : allocation_source
+    ; callstack : Printexc.raw_backtrace
+    }
+
+  type ('major, 'minor) tracker =
+    { alloc_minor : allocation -> 'minor option
+    ; alloc_major : allocation -> 'major option
+    ; promote : 'minor -> 'major option
+    ; dealloc_minor : 'minor -> unit
+    }
+
+  let null_tracker =
+    { alloc_minor = (fun _ -> None)
+    ; alloc_major = (fun _ -> None)
+    ; promote = (fun _ -> None)
+    ; dealloc_minor = (fun _ -> ())
+    }
+  ;;
+
+  let start ~sampling_rate:_ ?callstack_size:_ (_ : (_, _) tracker) = ()
+end
+
+[%%if ocaml_version >= (5, 4, 0)]
+
+open Gc
+
+[%%else]
+
+(* The [alloc] field is never mutated if it's disabled *)
+[@@@ocaml.warning "-69"]
+
+[%%endif]
+
+let sampling_rate = 1e-4
+let callstack_size = 10
+let top_entry_count = 10
+
+module Trace = struct
+  type t =
+    | Unknown
+    | Trace of Printexc.raw_backtrace_slot array
+
+  let equal t t' =
+    match t, t' with
+    | Unknown, Unknown -> true
+    | Trace slots, Trace slots' ->
+      let len = Array.length slots in
+      Int.equal len (Array.length slots')
+      &&
+      let rec loop i = i = len || (Poly.equal slots.(i) slots'.(i) && loop (i + 1)) in
+      loop 0
+    | Unknown, Trace _ | Trace _, Unknown -> false
+  ;;
+
+  let hash = function
+    | Unknown -> 0
+    | Trace slots ->
+      Array.fold_left slots ~init:(Array.length slots) ~f:(fun acc slot ->
+        acc * 65599 lxor Poly.hash slot)
+  ;;
+
+  let to_dyn = function
+    | Unknown -> Dyn.string "<unknown>"
+    | Trace _ -> Dyn.opaque ()
+  ;;
+end
+
+module Trace_table = Hashtbl.Make (Trace)
+
+type tracked_minor =
+  { trace : Trace.t
+  ; n_samples : int
+  }
+
+type heap =
+  { mutable total_samples : int
+  ; mutable by_trace : int Trace_table.t
+  }
+
+type t =
+  { mutex : Mutex.t
+  ; minor : heap
+  ; major : heap
+  ; promoted : heap
+  ; mutable profile : Memprof.t option
+  }
+
+let create_heap () = { total_samples = 0; by_trace = Trace_table.create 64 }
+
+let create () =
+  { mutex = Mutex.create ()
+  ; minor = create_heap ()
+  ; major = create_heap ()
+  ; promoted = create_heap ()
+  ; profile = None
+  }
+;;
+
+let place_of_slot slot =
+  let name = Printexc.Slot.name slot in
+  match Printexc.Slot.location slot with
+  | None -> Option.value name ~default:"<unknown>"
+  | Some { filename; line_number; start_char; _ } ->
+    (match name with
+     | None -> sprintf "%s:%d:%d" filename line_number start_char
+     | Some name -> sprintf "%s:%d:%d %s" filename line_number start_char name)
+;;
+
+let trace_to_strings = function
+  | Trace.Unknown -> [ "<unknown>" ]
+  | Trace slots ->
+    Array.to_list slots
+    |> List.map ~f:(fun slot -> Printexc.convert_raw_backtrace_slot slot |> place_of_slot)
+;;
+
+let trace_of_callstack callstack =
+  let length = Printexc.raw_backtrace_length callstack in
+  let rec add_inlined slot remaining acc =
+    if remaining = 0
+    then acc, remaining
+    else (
+      let acc, remaining = slot :: acc, remaining - 1 in
+      if remaining = 0
+      then acc, remaining
+      else (
+        match Printexc.get_raw_backtrace_next_slot slot with
+        | None -> acc, remaining
+        | Some slot -> add_inlined slot remaining acc))
+  in
+  let rec loop i remaining acc =
+    if i = length || remaining = 0
+    then (
+      match List.rev acc with
+      | [] -> Trace.Unknown
+      | slots -> Trace (Array.of_list slots))
+    else (
+      let acc, remaining =
+        add_inlined (Printexc.get_raw_backtrace_slot callstack i) remaining acc
+      in
+      loop (i + 1) remaining acc)
+  in
+  loop 0 callstack_size []
+;;
+
+let record_sample t heap ~trace ~n_samples =
+  Mutex.protect t.mutex (fun () ->
+    heap.total_samples <- heap.total_samples + n_samples;
+    match Trace_table.find heap.by_trace trace with
+    | None -> Trace_table.set heap.by_trace trace n_samples
+    | Some samples -> Trace_table.set heap.by_trace trace (samples + n_samples))
+;;
+
+let tracker t =
+  { Memprof.null_tracker with
+    alloc_minor =
+      (fun { Memprof.n_samples; callstack; _ } ->
+        let trace = trace_of_callstack callstack in
+        record_sample t t.minor ~trace ~n_samples;
+        Some { trace; n_samples })
+  ; alloc_major =
+      (fun { Memprof.n_samples; callstack; _ } ->
+        let trace = trace_of_callstack callstack in
+        record_sample t t.major ~trace ~n_samples;
+        None)
+  ; promote =
+      (fun { trace; n_samples } ->
+        record_sample t t.promoted ~trace ~n_samples;
+        None)
+  }
+;;
+
+let start () =
+  let t = create () in
+  let profile = Memprof.start ~sampling_rate ~callstack_size (tracker t) in
+  t.profile <- Some profile;
+  t
+;;
+
+[%%if ocaml_version >= (5, 4, 0)]
+
+let stop t =
+  Option.iter t.profile ~f:(fun profile ->
+    Memprof.stop ();
+    Memprof.discard profile;
+    t.profile <- None)
+;;
+
+[%%else]
+
+let stop t = t.profile <- None
+
+[%%endif]
+
+let estimated_words_of_samples samples =
+  int_of_float ((float_of_int samples /. sampling_rate) +. 0.5)
+;;
+
+let take_top_entries entries =
+  let rec take acc n = function
+    | _ when n <= 0 -> List.rev acc
+    | [] -> List.rev acc
+    | x :: xs -> take (x :: acc) (n - 1) xs
+  in
+  take [] top_entry_count entries
+;;
+
+let insert_top_entry entry entries =
+  let _, samples = entry in
+  let rec insert = function
+    | [] -> [ entry ]
+    | ((_, samples') as entry') :: entries ->
+      if samples > samples' then entry :: entry' :: entries else entry' :: insert entries
+  in
+  insert entries |> take_top_entries
+;;
+
+let top_entries by_trace =
+  Trace_table.foldi by_trace ~init:[] ~f:(fun trace samples entries ->
+    insert_top_entry (trace, samples) entries)
+  |> List.map ~f:(fun (trace, samples) ->
+    let estimated_words = estimated_words_of_samples samples in
+    ({ trace = trace_to_strings trace; estimated_words; samples } : Event.alloc_entry))
+;;
+
+let summary_of_heap total_samples by_trace =
+  let total_words = estimated_words_of_samples total_samples in
+  ({ total_words; total_samples; top = top_entries by_trace } : Event.alloc_heap)
+;;
+
+let swap t =
+  let fresh_minor = Trace_table.create 64 in
+  let fresh_major = Trace_table.create 64 in
+  let fresh_promoted = Trace_table.create 64 in
+  Mutex.protect t.mutex (fun () ->
+    let minor_total_samples = t.minor.total_samples in
+    let minor_by_trace = t.minor.by_trace in
+    let major_total_samples = t.major.total_samples in
+    let major_by_trace = t.major.by_trace in
+    let promoted_total_samples = t.promoted.total_samples in
+    let promoted_by_trace = t.promoted.by_trace in
+    t.minor.total_samples <- 0;
+    t.minor.by_trace <- fresh_minor;
+    t.major.total_samples <- 0;
+    t.major.by_trace <- fresh_major;
+    t.promoted.total_samples <- 0;
+    t.promoted.by_trace <- fresh_promoted;
+    ( minor_total_samples
+    , minor_by_trace
+    , major_total_samples
+    , major_by_trace
+    , promoted_total_samples
+    , promoted_by_trace ))
+;;
+
+type snapshot =
+  { minor : Event.alloc_heap
+  ; major : Event.alloc_heap
+  ; promoted : Event.alloc_heap
+  }
+
+let snapshot t =
+  let ( minor_total_samples
+      , minor_by_trace
+      , major_total_samples
+      , major_by_trace
+      , promoted_total_samples
+      , promoted_by_trace )
+    =
+    swap t
+  in
+  let minor = summary_of_heap minor_total_samples minor_by_trace in
+  let major = summary_of_heap major_total_samples major_by_trace in
+  let promoted = summary_of_heap promoted_total_samples promoted_by_trace in
+  { minor; major; promoted }
+;;
+
+let reset t =
+  ignore
+    (swap t : int * int Trace_table.t * int * int Trace_table.t * int * int Trace_table.t)
+;;

src/dune_trace/alloc.mli

@@ -0,0 +1,12 @@
+type t
+
+type snapshot =
+  { minor : Event.alloc_heap
+  ; major : Event.alloc_heap
+  ; promoted : Event.alloc_heap
+  }
+
+val start : unit -> t
+val reset : t -> unit
+val stop : t -> unit
+val snapshot : t -> snapshot