Multicore statmemprof support, with a caller-supplied Memprof implementation

README.md

@@ -23,3 +23,47 @@ command-line tools in bin/, but the recommended interface is the
 memtrace viewer, which lives at:
 
     https://github.com/janestreet/memtrace_viewer
+
+## Installation
+These instructions are for using statmemprof with OCaml 5.3.0+trunk
+
+``` shell
+# Setup a new Blank switch
+opam switch create 5.3.0 --no-install
+eval $(opam env --switch=5.3.0 --set-switch)
+
+# Install dune
+opam install dune
+```
+
+Now we can get memory traces for programs, here is a multicore fibonacci program:
+
+``` shell
+$ opam instal domainslib
+
+$ dune build examples
+
+# Run tracing on single domain
+$ MEMTRACE=fib_par.ctf _build/default/examples/fib_par.exe 1 45
+
+# On three domains
+$ MEMTRACE=fib_par_2.ctf _build/default/examples/fib_par.exe 3 45
+```
+
+these CTF files are viewable in `memtrace_viewer`.
+
+Install memtrace_viewer in another switch (5.1 will work since we only need to read and write trace files)
+
+``` shell
+opam switch create 5.1.1 --no-install
+opam install memtrace_viewer
+memtrace-viewer ./fib_par_2.ctf
+```
+
+
+## API changes
+
+- `start_tracing` and `trace_if_requested` take an optional
+  `?memprof:(module Memtrace.Memprof_sig)`, defaulting to `Gc.Memprof`.
+- `start_tracing` takes a final `()`.
+- Requires OCaml 5.3 or later.

docs/internal.md

@@ -118,16 +118,16 @@ arbitrary program is hard. There are three standard approaches,
 available as options in `perf record`:
 
   - `--call-graph=dwarf` uses the DWARF debugging information
-    
+
   - `--call-graph=fp` follows a chain of frame pointers
-  
+
   - `--call-graph=lbr` uses the Last Branch Record hardware support
 
 However, all of these have disadvantages:
 
   - DWARF exists to support debuggers, and so is designed for
     flexibility rather than speed.
-    
+
     This flexibility is necessary to handle the hard cases of C stack
     frames: for instance, a C program can define a variable-length
     array of ints on the stack, and store its length (in ints, not

dune

@@ -1 +1 @@
-(dirs bin src test trace_ocamlopt)
+(dirs bin src test trace_ocamlopt examples)

dune-project

@@ -14,4 +14,5 @@
  (synopsis "Streaming client for Memprof")
  (description "Generates compact traces of a program's memory use.")
  (depends
-  (ocaml (>= 4.11.0))))
+  (domainslib :with-test)
+  (ocaml (>= 5.3.0))))

examples/dune

@@ -0,0 +1,3 @@
+(executables
+ (libraries memtrace domainslib)
+ (names fib_par))

examples/fib_par.ml

@@ -0,0 +1,28 @@
+(* fib_par.ml *)
+let num_domains = try int_of_string Sys.argv.(1) with _ -> 1
+let n = try int_of_string Sys.argv.(2) with _ -> 1
+
+(* Sequential Fibonacci *)
+let rec fib n =
+  if n < 2 then 1 else fib (n - 1) + fib (n - 2)
+
+module T = Domainslib.Task
+
+let rec fib_par pool n =
+  let _ = Buffer.create 10000 in
+  if n > 20 then begin
+      let a = T.async pool (fun _ -> fib_par pool (n-1)) in
+      let b = T.async pool (fun _ -> fib_par pool (n-2)) in
+      T.await pool a + T.await pool b
+    end else
+    (* Call sequential Fibonacci if the available work is small *)
+    fib n
+
+let main () =
+  Memtrace.trace_if_requested ~context:"fib" ();
+  let pool = T.setup_pool ~num_domains:(num_domains - 1) () in
+  let res = T.run pool (fun _ -> fib_par pool n) in
+  T.teardown_pool pool;
+  Printf.printf "fib(%d) = %d\n" n res
+
+let _ = main ()

memtrace.opam

@@ -10,7 +10,8 @@ homepage: "https://github.com/janestreet/memtrace"
 bug-reports: "https://github.com/janestreet/memtrace/issues"
 depends: [
   "dune" {>= "2.3"}
-  "ocaml" {>= "4.11.0"}
+  "domainslib" {with-test}
+  "ocaml" {>= "5.3.0"}
 ]
 build: [
   ["dune" "subst"] {pinned}

src/memprof_tracer.ml

@@ -1,53 +1,49 @@
+module type Memprof_sig = sig
+  include module type of Stdlib.Gc.Memprof
+end
+
 type t =
-  { mutable locked : bool;
-    mutable locked_ext : bool;
-    mutable failed : bool;
-    mutable stopped : bool;
+  { failed : bool Atomic.t;
+    stopped : bool Atomic.t;
+    mutex : Mutex.t;
+    stop_memprof : unit -> unit;
     report_exn : exn -> unit;
     trace : Trace.Writer.t;
     ext_sampler : Geometric_sampler.t; }
 
-let curr_active_tracer : t option ref = ref None
+let curr_active_tracer : t option Atomic.t  = Atomic.make None
 
-let active_tracer () = !curr_active_tracer
+let active_tracer () = Atomic.get curr_active_tracer
 
-let bytes_before_ext_sample = ref max_int
+let bytes_before_ext_sample = Atomic.make max_int
 
 let draw_sampler_bytes t =
   Geometric_sampler.draw t.ext_sampler * (Sys.word_size / 8)
 
-let[@inline never] rec lock_tracer s =
-  if s.locked then
-    if s.locked_ext then false
-    else (Thread.yield (); lock_tracer s)
-  else if s.failed then
-    false
-  else
-    (s.locked <- true; true)
-
-let[@inline never] rec lock_tracer_ext s =
-  if s.locked then
-    (Thread.yield (); lock_tracer_ext s)
-  else if s.failed then
+let[@inline never] lock_tracer s =
+  if Atomic.get s.failed then
     false
-  else
-    (s.locked <- true; s.locked_ext <- true; true)
+    (* During external allocations or closing, a thread may try to obtain
+       a lock it already holds. In that case, Mutex.lock will throw
+       an error and we can ignore it. *)
+  else begin
+    try
+      Mutex.lock s.mutex;
+      (* The failed flag can be set while we wait for the mutex, so test it
+         again once we hold the lock. *)
+      if Atomic.get s.failed then (Mutex.unlock s.mutex; false) else true
+    with
+      | Sys_error _ -> false
+  end
 
 let[@inline never] unlock_tracer s =
-  assert (s.locked && not s.locked_ext && not s.failed);
-  s.locked <- false
-
-let[@inline never] unlock_tracer_ext s =
-  assert (s.locked && s.locked_ext && not s.failed);
-  s.locked_ext <- false;
-  s.locked <- false
+  assert (not (Atomic.get s.failed));
+  Mutex.unlock s.mutex
 
 let[@inline never] mark_failed s e =
-  assert (s.locked && not s.failed);
-  s.failed <- true;
-  s.locked <- false;
-  s.locked_ext <- false;
-  s.report_exn e
+  if (Atomic.compare_and_set s.failed false true) then
+    s.report_exn e;
+  Mutex.unlock s.mutex
 
 let default_report_exn e =
   match e with
@@ -61,9 +57,16 @@ let default_report_exn e =
      Printexc.print_backtrace stderr;
      flush stderr
 
-let start ?(report_exn=default_report_exn) ~sampling_rate trace =
+let default_memprof = (module Stdlib.Gc.Memprof : Memprof_sig)
+
+let start ?(report_exn=default_report_exn) ?(memprof = default_memprof)
+      ~sampling_rate trace =
+  let (module Memprof : Memprof_sig) = memprof in
   let ext_sampler = Geometric_sampler.make ~sampling_rate () in
-  let s = { trace; locked = false; locked_ext = false; stopped = false; failed = false;
+  let mutex = Mutex.create () in
+  let profile : Memprof.t option ref = ref None in
+  let s = { trace; mutex; stopped = Atomic.make false; failed = Atomic.make false;
+            stop_memprof = (fun () -> Memprof.stop (); Option.iter Memprof.discard !profile);
             report_exn; ext_sampler } in
   let tracker : (_,_) Gc.Memprof.tracker = {
     alloc_minor = (fun info ->
@@ -75,8 +78,11 @@ let start ?(report_exn=default_report_exn) ~sampling_rate trace =
                 ~callstack:info.callstack
         with
         | r -> unlock_tracer s; Some r
-        | exception e -> mark_failed s e; None
-      end else None);
+        | exception e ->
+           mark_failed s e;
+           None
+        end
+      else None);
     alloc_major = (fun info ->
       if lock_tracer s then begin
         match Trace.Writer.put_alloc_with_raw_backtrace trace (Trace.Timestamp.now ())
@@ -104,37 +110,39 @@ let start ?(report_exn=default_report_exn) ~sampling_rate trace =
         match Trace.Writer.put_collect trace (Trace.Timestamp.now ()) id with
         | () -> unlock_tracer s
         | exception e -> mark_failed s e) } in
-  curr_active_tracer := Some s;
-  bytes_before_ext_sample := draw_sampler_bytes s;
-  Gc.Memprof.start
-    ~sampling_rate
-    ~callstack_size:max_int
-    tracker;
+  Atomic.set curr_active_tracer (Some s);
+  Atomic.set bytes_before_ext_sample (draw_sampler_bytes s);
+  profile := Some (Memprof.start ~sampling_rate ~callstack_size:max_int tracker);
   s
 
 let stop s =
-  if not s.stopped then begin
-    s.stopped <- true;
-    Gc.Memprof.stop ();
+  (* Call stop to stop sampling on the current profile.
+     Promotion and deallocation callbacks from a profile may run
+     after stop is called, however we ignore these callbacks when
+     stopping.
+   *)
+  if (Atomic.compare_and_set s.stopped false true) then begin
+    s.stop_memprof ();
     if lock_tracer s then begin
-      try Trace.Writer.close s.trace with e -> mark_failed s e
+      try Trace.Writer.close s.trace with e ->
+        (Atomic.set s.failed true; s.report_exn e);
+      Mutex.unlock s.mutex
     end;
-    curr_active_tracer := None
+    Atomic.set curr_active_tracer None
   end
 
 let[@inline never] ext_alloc_slowpath ~bytes =
-  match !curr_active_tracer with
-  | None -> bytes_before_ext_sample := max_int; None
+  match Atomic.get curr_active_tracer with
+  | None -> Atomic.set bytes_before_ext_sample max_int; None
   | Some s ->
-    if lock_tracer_ext s then begin
+    if lock_tracer s then begin
       match
         let bytes_per_word = Sys.word_size / 8 in
         (* round up to an integer number of words *)
         let size_words = (bytes + bytes_per_word - 1) / bytes_per_word in
         let samples = ref 0 in
-        while !bytes_before_ext_sample <= 0 do
-          bytes_before_ext_sample :=
-            !bytes_before_ext_sample + draw_sampler_bytes s;
+        while Atomic.get bytes_before_ext_sample <= 0 do
+          ignore (Atomic.fetch_and_add bytes_before_ext_sample (draw_sampler_bytes s));
           incr samples
         done;
         assert (!samples > 0);
@@ -146,26 +154,24 @@ let[@inline never] ext_alloc_slowpath ~bytes =
                 ~source:External
                 ~callstack)
       with
-      | r -> unlock_tracer_ext s; r
+      | r -> unlock_tracer s; r
       | exception e -> mark_failed s e; None
     end else None
 
-
 type ext_token = Trace.Obj_id.t
 
 let ext_alloc ~bytes =
-  let n = !bytes_before_ext_sample - bytes in
-  bytes_before_ext_sample := n;
+  let n = Atomic.fetch_and_add bytes_before_ext_sample (- bytes) - bytes in
   if n <= 0 then ext_alloc_slowpath ~bytes else None
 
 let ext_free id =
-  match !curr_active_tracer with
+  match Atomic.get curr_active_tracer with
   | None -> ()
   | Some s ->
-    if lock_tracer_ext s then begin
+    if lock_tracer s then begin
       match
         Trace.Writer.put_collect s.trace (Trace.Timestamp.now ()) id
       with
-      | () -> unlock_tracer_ext s; ()
+      | () -> unlock_tracer s; ()
       | exception e -> mark_failed s e; ()
     end

src/memprof_tracer.mli

@@ -1,5 +1,17 @@
+module type Memprof_sig = sig
+  include module type of Stdlib.Gc.Memprof
+end
+
 type t
-val start : ?report_exn:(exn -> unit) -> sampling_rate:float -> Trace.Writer.t -> t
+
+val default_memprof : (module Memprof_sig)
+
+val start :
+  ?report_exn:(exn -> unit) ->
+  ?memprof:(module Memprof_sig) ->
+  sampling_rate:float ->
+  Trace.Writer.t ->
+  t
 val stop : t -> unit
 
 val active_tracer : unit -> t option

src/memtrace.ml

@@ -2,7 +2,11 @@ type tracer = Memprof_tracer.t
 
 let getpid64 () = Int64.of_int (Unix.getpid ())
 
-let start_tracing ~context ~sampling_rate ~filename =
+module type Memprof_sig = Memprof_tracer.Memprof_sig
+
+let default_memprof = Memprof_tracer.default_memprof
+
+let start_tracing ?(memprof = default_memprof) ~context ~sampling_rate ~filename () =
   if Memprof_tracer.active_tracer () <> None then
     failwith "Only one Memtrace instance may be active at a time";
   let fd =
@@ -36,7 +40,7 @@ let start_tracing ~context ~sampling_rate ~filename =
       context;
     } in
   let trace = Trace.Writer.create fd ~getpid:getpid64 info in
-  Memprof_tracer.start ~sampling_rate trace
+  Memprof_tracer.start ~memprof ~sampling_rate trace
 
 let stop_tracing t =
   Memprof_tracer.stop t
@@ -46,7 +50,7 @@ let () =
 
 let default_sampling_rate = 1e-6
 
-let trace_if_requested ?context ?sampling_rate () =
+let trace_if_requested ?(memprof = default_memprof) ?context ?sampling_rate () =
   match Sys.getenv_opt "MEMTRACE" with
   | None | Some "" -> ()
   | Some filename ->
@@ -60,13 +64,13 @@ let trace_if_requested ?context ?sampling_rate () =
      in
      let sampling_rate =
        match Sys.getenv_opt "MEMTRACE_RATE" with
-       | Some rate -> check_rate (float_of_string_opt rate)
        | None | Some "" ->
-         match sampling_rate with
+         begin match sampling_rate with
          | Some _ -> check_rate sampling_rate
-         | None -> default_sampling_rate
+         | None -> default_sampling_rate end
+       | Some rate -> check_rate (float_of_string_opt rate)
      in
-     let _s = start_tracing ~context ~sampling_rate ~filename in
+     let _s = start_tracing ~memprof ~context ~sampling_rate ~filename () in
      ()
 
 module Trace = Trace