Merge pull request #10025 from stedolan/memprof-custom

Track caml_alloc_custom_mem allocations with Statmemprof

Changes

@@ -5,6 +5,10 @@ Working version
 
 ### Runtime system:
 
+- #10025: Track custom blocks (e.g. Bigarray) with Statmemprof
+  (Stephen Dolan, review by Leo White, Gabriel Scherer and Jacques-Henri
+   Jourdan)
+
 ### Code generation and optimizations:
 
 ### Standard library:

runtime/caml/memprof.h

@@ -27,6 +27,7 @@ extern void caml_memprof_set_suspended(int);
 extern value caml_memprof_handle_postponed_exn(void);
 
 extern void caml_memprof_track_alloc_shr(value block);
+extern void caml_memprof_track_custom(value block, mlsize_t bytes);
 extern void caml_memprof_track_young(uintnat wosize, int from_caml,
                                      int nallocs, unsigned char* alloc_lens);
 extern void caml_memprof_track_interned(header_t* block, header_t* blockend);

runtime/custom.c

@@ -24,6 +24,7 @@
 #include "caml/memory.h"
 #include "caml/mlvalues.h"
 #include "caml/signals.h"
+#include "caml/memprof.h"
 
 uintnat caml_custom_major_ratio = Custom_major_ratio_def;
 uintnat caml_custom_minor_ratio = Custom_minor_ratio_def;
@@ -102,7 +103,9 @@ CAMLexport value caml_alloc_custom_mem(struct custom_operations * ops,
     Bsize_wsize (Caml_state->stat_heap_wsz) / 150 * caml_custom_major_ratio;
   mlsize_t max_minor =
     Bsize_wsize (Caml_state->minor_heap_wsz) / 100 * caml_custom_minor_ratio;
-  return alloc_custom_gen (ops, bsz, mem, max_major, mem_minor, max_minor);
+  value v = alloc_custom_gen (ops, bsz, mem, max_major, mem_minor, max_minor);
+  caml_memprof_track_custom(v, mem);
+  return v;
 }
 
 struct custom_operations_list {

runtime/memprof.c

@@ -57,6 +57,9 @@ static intnat callstack_size;
 
 static value tracker;
 
+/* Gc.Memprof.allocation_source */
+enum { SRC_NORMAL = 0, SRC_MARSHAL = 1, SRC_CUSTOM = 2 };
+
 struct tracked {
   /* Memory block being sampled. This is a weak GC root. */
   value block;
@@ -79,8 +82,8 @@ struct tracked {
   /* Whether this block has been initially allocated in the minor heap. */
   unsigned int alloc_young : 1;
 
-  /* Whether this block comes from unmarshalling. */
-  unsigned int unmarshalled : 1;
+  /* The source of the allocation: normal allocations, marshal or custom_mem. */
+  unsigned int source : 2;
 
   /* Whether this block has been promoted. Implies [alloc_young]. */
   unsigned int promoted : 1;
@@ -405,7 +408,7 @@ static int realloc_entries(struct entry_array* ea, uintnat grow)
 #define Invalid_index (~(uintnat)0)
 
 Caml_inline uintnat new_tracked(uintnat n_samples, uintnat wosize,
-                                int is_unmarshalled, int is_young,
+                                int source, int is_young,
                                 value block, value user_data)
 {
   struct tracked *t;
@@ -419,7 +422,7 @@ Caml_inline uintnat new_tracked(uintnat n_samples, uintnat wosize,
   t->user_data = user_data;
   t->running = NULL;
   t->alloc_young = is_young;
-  t->unmarshalled = is_unmarshalled;
+  t->source = source;
   t->promoted = 0;
   t->deallocated = 0;
   t->cb_promote_called = t->cb_dealloc_called = 0;
@@ -513,7 +516,7 @@ static value run_alloc_callback_exn(uintnat t_idx)
   sample_info = caml_alloc_small(4, 0);
   Field(sample_info, 0) = Val_long(t->n_samples);
   Field(sample_info, 1) = Val_long(t->wosize);
-  Field(sample_info, 2) = Val_long(t->unmarshalled);
+  Field(sample_info, 2) = Val_long(t->source);
   Field(sample_info, 3) = t->user_data;
   return run_callback_exn(&local->entries, t_idx,
      t->alloc_young ? Alloc_minor(tracker) : Alloc_major(tracker), sample_info);
@@ -747,24 +750,37 @@ void caml_memprof_invert_tracked(void)
 
 /**** Sampling procedures ****/
 
-void caml_memprof_track_alloc_shr(value block)
+static void maybe_track_block(value block, uintnat n_samples,
+                              uintnat wosize, int src)
 {
-  uintnat n_samples;
-  value callstack = 0;
-  CAMLassert(Is_in_heap(block));
-
-  if (lambda == 0 || local->suspended) return;
-
-  n_samples = rand_binom(Whsize_val(block));
+  value callstack;
   if (n_samples == 0) return;
 
   callstack = capture_callstack_postponed();
   if (callstack == 0) return;
 
-  new_tracked(n_samples, Wosize_val(block), 0, 0, block, callstack);
+  new_tracked(n_samples, wosize, src, Is_young(block), block, callstack);
   check_action_pending();
 }
 
+void caml_memprof_track_alloc_shr(value block)
+{
+  CAMLassert(Is_in_heap(block));
+  if (lambda == 0 || local->suspended) return;
+
+  maybe_track_block(block, rand_binom(Whsize_val(block)),
+                    Wosize_val(block), SRC_NORMAL);
+}
+
+void caml_memprof_track_custom(value block, mlsize_t bytes)
+{
+  CAMLassert(Is_young(block) || Is_in_heap(block));
+  if (lambda == 0 || local->suspended) return;
+
+  maybe_track_block(block, rand_binom(Wsize_bsize(bytes)),
+                    Wsize_bsize(bytes), SRC_CUSTOM);
+}
+
 /* Shifts the next sample in the minor heap by [n] words. Essentially,
    this tells the sampler to ignore the next [n] words of the minor
    heap. */
@@ -822,13 +838,8 @@ void caml_memprof_track_young(uintnat wosize, int from_caml,
       rand_binom(caml_memprof_young_trigger - 1 - Caml_state->young_ptr);
     CAMLassert(encoded_alloc_lens == NULL);    /* No Comballoc in C! */
     caml_memprof_renew_minor_sample();
-
-    callstack = capture_callstack_postponed();
-    if (callstack == 0) return;
-
-    new_tracked(n_samples, wosize,
-                0, 1, Val_hp(Caml_state->young_ptr), callstack);
-    check_action_pending();
+    maybe_track_block(Val_hp(Caml_state->young_ptr), n_samples,
+                      wosize, SRC_NORMAL);
     return;
   }
 
@@ -866,8 +877,8 @@ void caml_memprof_track_young(uintnat wosize, int from_caml,
       int stopped;
 
       callstack = capture_callstack(alloc_idx);
-      t_idx = new_tracked(n_samples, alloc_wosz,
-                          0, 1, Placeholder_offs(alloc_ofs), callstack);
+      t_idx = new_tracked(n_samples, alloc_wosz, SRC_NORMAL, 1,
+                          Placeholder_offs(alloc_ofs), callstack);
       if (t_idx == Invalid_index) continue;
       res = run_alloc_callback_exn(t_idx);
       /* Has [caml_memprof_stop] been called during the callback? */
@@ -987,7 +998,7 @@ void caml_memprof_track_interned(header_t* block, header_t* blockend)
     if (callstack == 0) callstack = capture_callstack_postponed();
     if (callstack == 0) break;  /* OOM */
     new_tracked(rand_binom(next_p - next_sample_p) + 1,
-                Wosize_hp(p), 1, is_young, Val_hp(p), callstack);
+                Wosize_hp(p), SRC_MARSHAL, is_young, Val_hp(p), callstack);
     p = next_p;
   }
   check_action_pending();

stdlib/gc.ml

@@ -125,10 +125,11 @@ let delete_alarm a = a := false
 
 module Memprof =
   struct
+    type allocation_source = Normal | Marshal | Custom
     type allocation =
       { n_samples : int;
         size : int;
-        unmarshalled : bool;
+        source : allocation_source;
         callstack : Printexc.raw_backtrace }
 
     type ('minor, 'major) tracker = {

stdlib/gc.mli

@@ -463,6 +463,7 @@ external eventlog_resume : unit -> unit = "caml_eventlog_resume"
    notice. *)
 module Memprof :
   sig
+    type allocation_source = Normal | Marshal | Custom
     type allocation = private
       { n_samples : int;
         (** The number of samples in this block (>= 1). *)
@@ -470,8 +471,8 @@ module Memprof :
         size : int;
         (** The size of the block, in words, excluding the header. *)
 
-        unmarshalled : bool;
-        (** Whether the block comes from unmarshalling. *)
+        source : allocation_source;
+        (** The type of the allocation. *)
 
         callstack : Printexc.raw_backtrace
         (** The callstack for the allocation. *)

testsuite/tests/statmemprof/arrays_in_major.ml

@@ -112,7 +112,7 @@ let check_distrib lo hi cnt rate =
       alloc_major = (fun info ->
         assert (info.size >= lo && info.size <= hi);
         assert (info.n_samples > 0);
-        assert (not info.unmarshalled);
+        assert (info.source = Normal);
         smp := !smp + info.n_samples;
         None
       );

testsuite/tests/statmemprof/arrays_in_minor.ml

@@ -126,7 +126,7 @@ let check_distrib lo hi cnt rate =
       alloc_minor = (fun info ->
         assert (info.size >= lo && info.size <= hi);
         assert (info.n_samples > 0);
-        assert (not info.unmarshalled);
+        assert (info.source = Normal);
         smp := !smp + info.n_samples;
         None
       );

testsuite/tests/statmemprof/custom.ml

@@ -0,0 +1,44 @@
+(* TEST *)
+
+open Gc.Memprof
+
+let bigstring_create sz =
+  Bigarray.Array1.create Bigarray.char Bigarray.c_layout sz
+
+let keep = ref []
+
+let test sampling_rate =
+  let size = 256 in
+  let iters = 100_000 in
+  let size_words = size / (Sys.word_size / 8) in
+  let alloc = ref 0 and collect = ref 0 and promote = ref 0 in
+  let tracker =
+    { null_tracker with
+      alloc_minor = (fun info ->
+        if info.source <> Custom then None
+        else begin
+          alloc := !alloc + info.n_samples;
+          Some info.n_samples
+        end);
+      promote = (fun ns ->
+        promote := !promote + ns; None);
+      dealloc_minor = (fun ns ->
+        collect := !collect + ns) } in
+  start ~sampling_rate tracker;
+  for i = 1 to iters do
+    let str = Sys.opaque_identity bigstring_create size in
+    if i mod 10 = 0 then keep := str :: !keep
+  done;
+  keep := [];
+  Gc.full_major ();
+  stop ();
+  assert (!alloc = !promote + !collect);
+  let iters = float_of_int iters and size_words = float_of_int size_words in
+  (* see comballoc.ml for notes on precision *)
+  Printf.printf "%.2f %.1f\n"
+    ((float_of_int !alloc /. iters) /. size_words)
+    ((float_of_int !promote /. iters) /. size_words *. 10.)
+
+
+let () =
+  [0.01; 0.5; 0.17] |> List.iter test

testsuite/tests/statmemprof/custom.reference

@@ -0,0 +1,3 @@
+0.01 0.0
+0.50 0.5
+0.17 0.2

testsuite/tests/statmemprof/intern.ml

@@ -134,7 +134,7 @@ let check_distrib lo hi cnt rate =
   let alloc info =
     (* We also allocate the list constructor in the minor heap,
        so we filter that out. *)
-    if info.unmarshalled then begin
+    if info.source = Marshal then begin
       assert (info.size = 1 || info.size = 2);
       assert (info.n_samples > 0);
       smp := !smp + info.n_samples

testsuite/tests/statmemprof/lists_in_minor.ml

@@ -22,7 +22,7 @@ let check_distrib len cnt rate =
       alloc_minor = (fun info ->
         assert (info.size = 2);
         assert (info.n_samples > 0);
-        assert (not info.unmarshalled);
+        assert (info.source = Normal);
         smp := !smp + info.n_samples;
         None);
     };