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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
type t =
  { mutable locked : bool;
    mutable locked_ext : bool;
    mutable failed : bool;
    mutable stopped : bool;
    report_exn : exn -> unit;
    trace : Trace.Writer.t;
    ext_sampler : Geometric_sampler.t; }

let curr_active_tracer : t option ref = ref None

let active_tracer () = !curr_active_tracer

let bytes_before_ext_sample = ref 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
    false
  else
    (s.locked <- true; s.locked_ext <- true; true)

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

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

let default_report_exn e =
  match e with
  | Trace.Writer.Pid_changed ->
     (* This error is silently ignored, so that if Memtrace is active across
        Unix.fork () then the child process silently stops tracing *)
     ()
  | 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 ext_sampler = Geometric_sampler.make ~sampling_rate () in
  let s = { trace; locked = false; locked_ext = false; stopped = false; failed = false;
            report_exn; ext_sampler } 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
                ~source:Minor
                ~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
                ~source:Major
                ~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
  curr_active_tracer := Some s;
  bytes_before_ext_sample := draw_sampler_bytes s;
  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 begin
      try Trace.Writer.close s.trace with e -> mark_failed s e
    end;
    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
  | Some s ->
    if lock_tracer_ext 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;
          incr samples
        done;
        assert (!samples > 0);
        let callstack = Printexc.get_callstack max_int in
        Some (Trace.Writer.put_alloc_with_raw_backtrace s.trace
                (Trace.Timestamp.now ())
                ~length:size_words
                ~nsamples:!samples
                ~source:External
                ~callstack)
      with
      | r -> unlock_tracer_ext 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;
  if n <= 0 then ext_alloc_slowpath ~bytes else None

let ext_free id =
  match !curr_active_tracer with
  | None -> ()
  | Some s ->
    if lock_tracer_ext s then begin
      match
        Trace.Writer.put_collect s.trace (Trace.Timestamp.now ()) id
      with
      | () -> unlock_tracer_ext s; ()
      | exception e -> mark_failed s e; ()
    end