Source file `memprof_tracer.ml`

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
type t =
  { mutable locked : bool;
    mutable failed : bool;
    mutable stopped : bool;
    report_exn : exn -> unit;
    trace : Trace.Writer.t }

let[@inline never] lock_tracer s =
  (* This is a maximally unfair spinlock. *)
  (* if s.locked then Printf.fprintf stderr "contention\n%!"; *)
  while s.locked do Thread.yield () done;
  if s.failed then
    false
  else
    (s.locked <- true; true)

let[@inline never] unlock_tracer s =
  assert (s.locked && not s.failed);
  s.locked <- false

let[@inline never] mark_failed s e =
  assert (s.locked && not s.failed);
  s.failed <- true;
  s.locked <- false;
  s.report_exn e

let default_report_exn e =
  let msg = Printf.sprintf "Memtrace failure: %s\n" (Printexc.to_string e) in
  output_string stderr msg;
  Printexc.print_backtrace stderr;
  flush stderr

let start ?(report_exn=default_report_exn) ~sampling_rate trace =
  let s = { trace; locked = false; stopped = false; failed = false; report_exn } in
  let tracker : (_,_) Gc.Memprof.tracker = {
    alloc_minor = (fun info ->
      if lock_tracer s then begin
        match Trace.Writer.put_alloc_with_raw_backtrace trace (Trace.Timestamp.now ())
                ~length:info.size
                ~nsamples:info.n_samples
                ~is_major:false
                ~callstack:info.callstack
        with
        | r -> unlock_tracer s; Some r
        | 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 ())
                ~length:info.size
                ~nsamples:info.n_samples
                ~is_major:true
                ~callstack:info.callstack
        with
        | r -> unlock_tracer s; Some r
        | exception e -> mark_failed s e; None
      end else None);
    promote = (fun id ->
      if lock_tracer s then
        match Trace.Writer.put_promote trace (Trace.Timestamp.now ()) id with
        | () -> unlock_tracer s; Some id
        | exception e -> mark_failed s e; None
      else None);
    dealloc_minor = (fun id ->
      if lock_tracer s then
        match Trace.Writer.put_collect trace (Trace.Timestamp.now ()) id with
        | () -> unlock_tracer s
        | exception e -> mark_failed s e);
    dealloc_major = (fun id ->
      if lock_tracer s then
        match Trace.Writer.put_collect trace (Trace.Timestamp.now ()) id with
        | () -> unlock_tracer s
        | exception e -> mark_failed s e) } in
  Gc.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 ();
    if lock_tracer s then
      Trace.Writer.close s.trace
  end