Track allocation sources in allocation traces

src/dune_trace/alloc.ml

@@ -6,6 +6,10 @@ module Memprof = struct
   type t = unit
   type allocation_source = |
 
+  let string_of_allocation_source : allocation_source -> string = function
+    | _ -> .
+  ;;
+
   type allocation =
     { n_samples : int
     ; size : int
@@ -76,16 +80,37 @@ module Trace = struct
   ;;
 end
 
-module Trace_table = Hashtbl.Make (Trace)
+module Source = struct
+  type t = Memprof.allocation_source
+
+  let equal = Poly.equal
+  let hash = Poly.hash
+  let to_string = Memprof.string_of_allocation_source
+  let to_dyn t = Dyn.string (to_string t)
+end
+
+module Key = struct
+  type t =
+    { source : Source.t
+    ; trace : Trace.t
+    }
+
+  let equal x y = Source.equal x.source y.source && Trace.equal x.trace y.trace
+  let hash { source; trace } = Source.hash source * 65599 lxor Trace.hash trace
+
+  let to_dyn { source; trace } =
+    Dyn.record [ "source", Source.to_dyn source; "trace", Trace.to_dyn trace ]
+  ;;
+end
 
 type tracked_minor =
-  { trace : Trace.t
+  { key : Key.t
   ; n_samples : int
   }
 
 type heap =
   { mutable total_samples : int
-  ; mutable by_trace : int Trace_table.t
+  ; mutable by_key : (Key.t, int) Table.t
   }
 
 type t =
@@ -96,7 +121,7 @@ type t =
   ; mutable profile : Memprof.t option
   }
 
-let create_heap () = { total_samples = 0; by_trace = Trace_table.create 64 }
+let create_heap () = { total_samples = 0; by_key = Table.create (module Key) 64 }
 
 let create () =
   { mutex = Mutex.create ()
@@ -153,29 +178,34 @@ let trace_of_callstack callstack =
   loop 0 callstack_size []
 ;;
 
-let record_sample t heap ~trace ~n_samples =
+let record_sample t heap ~key ~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))
+    match Table.find heap.by_key key with
+    | None -> Table.set heap.by_key key n_samples
+    | Some samples -> Table.set heap.by_key key (samples + n_samples))
+;;
+
+let key_of_allocation { Memprof.source; callstack; _ } =
+  let trace = trace_of_callstack callstack in
+  { Key.source; trace }
 ;;
 
 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 })
+      (fun ({ Memprof.n_samples; _ } as allocation) ->
+        let key = key_of_allocation allocation in
+        record_sample t t.minor ~key ~n_samples;
+        Some { key; n_samples })
   ; alloc_major =
-      (fun { Memprof.n_samples; callstack; _ } ->
-        let trace = trace_of_callstack callstack in
-        record_sample t t.major ~trace ~n_samples;
+      (fun ({ Memprof.n_samples; _ } as allocation) ->
+        let key = key_of_allocation allocation in
+        record_sample t t.major ~key ~n_samples;
         None)
   ; promote =
-      (fun { trace; n_samples } ->
-        record_sample t t.promoted ~trace ~n_samples;
+      (fun { key; n_samples } ->
+        record_sample t t.promoted ~key ~n_samples;
         None)
   }
 ;;
@@ -225,42 +255,68 @@ let insert_top_entry entry entries =
   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 top_entries by_key =
+  Table.foldi by_key ~init:[] ~f:(fun key samples entries ->
+    insert_top_entry (key, samples) entries)
+  |> List.map ~f:(fun ({ Key.source; trace }, samples) ->
+    let estimated_words = estimated_words_of_samples samples in
+    ({ source = Source.to_string source
+     ; trace = trace_to_strings trace
+     ; estimated_words
+     ; samples
+     }
+     : Event.alloc_entry))
+;;
+
+let source_entries by_key =
+  let by_source = Table.create (module Source) 4 in
+  Table.foldi by_key ~init:() ~f:(fun { Key.source; _ } samples () ->
+    match Table.find by_source source with
+    | None -> Table.set by_source source samples
+    | Some previous -> Table.set by_source source (previous + samples));
+  Table.to_list by_source
+  |> List.sort ~compare:(fun (source, samples) (source', samples') ->
+    match Int.compare samples' samples with
+    | Eq -> String.compare (Source.to_string source) (Source.to_string source')
+    | ordering -> ordering)
+  |> List.map ~f:(fun (source, samples) ->
     let estimated_words = estimated_words_of_samples samples in
-    ({ trace = trace_to_strings trace; estimated_words; samples } : Event.alloc_entry))
+    ({ source = Source.to_string source; estimated_words; samples } : Event.alloc_source))
 ;;
 
-let summary_of_heap total_samples by_trace =
+let summary_of_heap total_samples by_key =
   let total_words = estimated_words_of_samples total_samples in
-  ({ total_words; total_samples; top = top_entries by_trace } : Event.alloc_heap)
+  ({ total_words
+   ; total_samples
+   ; by_source = source_entries by_key
+   ; top = top_entries by_key
+   }
+   : 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
+  let fresh_minor = Table.create (module Key) 64 in
+  let fresh_major = Table.create (module Key) 64 in
+  let fresh_promoted = Table.create (module Key) 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 minor_by_key = t.minor.by_key in
     let major_total_samples = t.major.total_samples in
-    let major_by_trace = t.major.by_trace in
+    let major_by_key = t.major.by_key in
     let promoted_total_samples = t.promoted.total_samples in
-    let promoted_by_trace = t.promoted.by_trace in
+    let promoted_by_key = t.promoted.by_key in
     t.minor.total_samples <- 0;
-    t.minor.by_trace <- fresh_minor;
+    t.minor.by_key <- fresh_minor;
     t.major.total_samples <- 0;
-    t.major.by_trace <- fresh_major;
+    t.major.by_key <- fresh_major;
     t.promoted.total_samples <- 0;
-    t.promoted.by_trace <- fresh_promoted;
+    t.promoted.by_key <- fresh_promoted;
     ( minor_total_samples
-    , minor_by_trace
+    , minor_by_key
     , major_total_samples
-    , major_by_trace
+    , major_by_key
     , promoted_total_samples
-    , promoted_by_trace ))
+    , promoted_by_key ))
 ;;
 
 type snapshot =
@@ -271,21 +327,27 @@ type snapshot =
 
 let snapshot t =
   let ( minor_total_samples
-      , minor_by_trace
+      , minor_by_key
       , major_total_samples
-      , major_by_trace
+      , major_by_key
       , promoted_total_samples
-      , promoted_by_trace )
+      , promoted_by_key )
     =
     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
+  let minor = summary_of_heap minor_total_samples minor_by_key in
+  let major = summary_of_heap major_total_samples major_by_key in
+  let promoted = summary_of_heap promoted_total_samples promoted_by_key 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)
+    (swap t
+     : int
+       * (Key.t, int) Table.t
+       * int
+       * (Key.t, int) Table.t
+       * int
+       * (Key.t, int) Table.t)
 ;;

src/dune_trace/dune_trace.mli

@@ -39,15 +39,23 @@ module Event : sig
 
   type t
 
+  type alloc_source =
+    { source : string
+    ; estimated_words : int
+    ; samples : int
+    }
+
   type alloc_entry =
-    { trace : string list
+    { source : string
+    ; trace : string list
     ; estimated_words : int
     ; samples : int
     }
 
   type alloc_heap =
     { total_words : int
     ; total_samples : int
+    ; by_source : alloc_source list
     ; top : alloc_entry list
     }
 

src/dune_trace/event.ml

@@ -111,15 +111,23 @@ end
 
 type t = Event.t
 
+type alloc_source =
+  { source : string
+  ; estimated_words : int
+  ; samples : int
+  }
+
 type alloc_entry =
-  { trace : string list
+  { source : string
+  ; trace : string list
   ; estimated_words : int
   ; samples : int
   }
 
 type alloc_heap =
   { total_words : int
   ; total_samples : int
+  ; by_source : alloc_source list
   ; top : alloc_entry list
   }
 
@@ -275,18 +283,28 @@ let watch_build_finish ~run_id ~outcome ~start ~stop ~restart_duration =
 
 let alloc_summary ~phase ~run_id ~minor ~major ~promoted =
   let now = Time.now () in
-  let entry { trace; estimated_words; samples } =
+  let source ({ source; estimated_words; samples } : alloc_source) =
+    Arg.record
+      [ "source", Arg.string source
+      ; "estimated_words", Arg.int estimated_words
+      ; "samples", Arg.int samples
+      ]
+    |> Arg.list
+  in
+  let entry ({ source; trace; estimated_words; samples } : alloc_entry) =
     Arg.record
-      [ "trace", Arg.list (List.map trace ~f:Arg.string)
+      [ "source", Arg.string source
+      ; "trace", Arg.list (List.map trace ~f:Arg.string)
       ; "estimated_words", Arg.int estimated_words
       ; "samples", Arg.int samples
       ]
     |> Arg.list
   in
-  let heap { total_words; total_samples; top } =
+  let heap { total_words; total_samples; by_source; top } =
     Arg.record
       [ "total_words", Arg.int total_words
       ; "total_samples", Arg.int total_samples
+      ; "by_source", Arg.list (List.map by_source ~f:source)
       ; "top", Arg.list (List.map top ~f:entry)
       ]
     |> Arg.list

test/blackbox-tests/test-cases/trace/alloc.t

@@ -24,12 +24,18 @@ The alloc sampler is only enabled when the alloc trace category is requested:
   >         , major: (.args.major | keys)
   >         , promoted: (.args.promoted | keys)
   >         }
-  >     , top_entries_are_traces:
+  >     , entries_have_sources_and_traces:
   >         (all((.args.minor.top + .args.major.top + .args.promoted.top)[]?;
-  >           ((keys | sort) == ["estimated_words", "samples", "trace"]
+  >           ((keys | sort) == ["estimated_words", "samples", "source", "trace"]
+  >            and (.source | type == "string")
   >            and (.trace | type == "array")
   >            and (.trace | length <= 10)
-  >            and all(.trace[]; type == "string"))))
+  >            and all(.trace[]; type == "string")))
+  >          and all((.args.minor.by_source
+  >                   + .args.major.by_source
+  >                   + .args.promoted.by_source)[]?;
+  >            ((keys | sort) == ["estimated_words", "samples", "source"]
+  >             and (.source | type == "string"))))
   >     }
   >   ]'
   [
@@ -39,44 +45,50 @@ The alloc sampler is only enabled when the alloc trace category is requested:
       "has_run_id": true,
       "heaps": {
         "minor": [
+          "by_source",
           "top",
           "total_samples",
           "total_words"
         ],
         "major": [
+          "by_source",
           "top",
           "total_samples",
           "total_words"
         ],
         "promoted": [
+          "by_source",
           "top",
           "total_samples",
           "total_words"
         ]
       },
-      "top_entries_are_traces": true
+      "entries_have_sources_and_traces": true
     },
     {
       "name": "summary",
       "phase": "exit",
       "has_run_id": false,
       "heaps": {
         "minor": [
+          "by_source",
           "top",
           "total_samples",
           "total_words"
         ],
         "major": [
+          "by_source",
           "top",
           "total_samples",
           "total_words"
         ],
         "promoted": [
+          "by_source",
           "top",
           "total_samples",
           "total_words"
         ]
       },
-      "top_entries_are_traces": true
+      "entries_have_sources_and_traces": true
     }
   ]