feat: switch hash to faster (#9)

Did some profiling of the profiler via Apple Instruments, our hashing takes more than 10% of the sampling time for large stack traces, so I switched us to collecting the parts in order!

Heres @jmgrosen explanation

> Sorry if I'm drawing this out longer than is necessary! But I'm worried about this global atomic ending up a bottleneck at high domain counts.
>
> After looking into the runtime events system a bit more, I don't think we even need a counter--each part just needs its bytes and two booleans `is_start` and `is_end`. Each domain gets its own ring buffer, so we don't have to worry about parts of different points being interleaved. The only things we need to do are  
> 1\. reassemble in-order parts of the same point  
> 2\. detect when a new point starts  
> 3\. detect when there has been a ring buffer overflow
>
> Assume for now there is only one domain/ring buffer. Then to implement `event_of_perf_event` we, starting at an `is_start` part, collect bytes in our buffer, until we hit an `is_end` part and return the unmarshalled form. This handles (1) and (2). To handle (3), we can use the `lost_events` callback from the runtime events system to put `event_of_perf_event` in a "throw parts away until the next `is_start` part" state. (The documentation is a little unclear, but based on the implementation, `lost_events` will _always_ be called when the ring buffer has overflowed.)
>
> Then for handling multiple ring buffers, we just store a different event buffer per-ring buffer, keyed on the `int` ring buffer ID that each of the runtime events callbacks receives.

## Test plan

compare runs of the example program with Apple Instruments, notice the hash function no longer shows up as major cpu usage.

Additionally, to test the logic of the new parts, I modified the example program to generate more events than necessary in order to overwrite the ring buffer leading to lost events, and ensured we did not segfault or anything.

Author: ajbt200128

dune-project

@@ -17,7 +17,7 @@
 (package
  (name pyro-caml-instruments)
  (synopsis "Pyroscope + OCaml = Pyro Caml")
- (depends ocaml logs digestif ppx_deriving)
+ (depends ocaml logs ppx_deriving)
  (tags ("profiling")))
 
 (package

lib/Event.ml

@@ -16,23 +16,18 @@
 (*****************************************************************************)
 (* Event *)
 (*****************************************************************************)
-(* md5 feels fast and safe enough *)
-module Hash = Digest.MD5
 
-(* an event is either a point or partial. We have a notion of a partial event
-   since runtime events only support payloads of max 1024 bytes, and any larger
-   will raise an exception. Since some callstacks can be LARGE!! we break up the
+(* runtime events only support payloads of max 1024 bytes, and any larger will
+   raise an exception. Since some callstacks can be LARGE!! we break up the
    callstack into a multipart message that is then reassembled by the
-   profiler *)
-type t =
-  | Point of (float * Stack_trace.raw_stack_trace)
-  | Partial of { id : Hash.t; bytes : Bytes.t; part : int; part_count : int }
+   profiler. *)
+type t = { bytes : Bytes.t; part : int; part_count : int }
 
+(* The actual underlying data we're transmitting, a stack trace with a
+   timestamp *)
+type point = float * Stack_trace.raw_stack_trace
 type marshaled = bytes * int
 
-let make_partial id part_count part bytes =
-  Partial { id; bytes; part; part_count }
-
 let split_bytes bytes size =
   let rec aux offset parts =
     if offset >= Bytes.length bytes then List.rev parts
@@ -44,27 +39,23 @@ let split_bytes bytes size =
   aux 0 []
 
 let marshal e =
-  let marshaled_event = Marshal.to_bytes e [] in
-  let len = Bytes.length marshaled_event in
-  (marshaled_event, len)
+  let marshaled_obj = Marshal.to_bytes e [] in
+  let len = Bytes.length marshaled_obj in
+  (marshaled_obj, len)
 
-(* 800 is chosen since when we break the event up into a partial event, we need
+(* 900 is chosen since when we break the event up into a partial event, we need
    to save some room for the other parts of the Partial data structure besides
    the bytes*)
 (* TODO be more clever about max size *)
-let marshal_event ?(max_size = 800) e =
-  let marshaled_event, len = marshal e in
-  (* Max size of runtime event type payload *)
+let marshal_point ?(max_size = 900) (p : point) =
+  let marshaled_point, _len = marshal p in
+  (* Max size of runtime event type payload is 1024, but we want to stay
+     slightly under that so we can store metadata about which part this is *)
   (* https://ocaml.org/manual/5.3/api/Runtime_events.Type.html *)
-  if len <= 1024 then [ (marshaled_event, len) ]
-  else
-    (* if it's bigger split it up! *)
-    let id = Hash.bytes marshaled_event in
-    let parts = split_bytes marshaled_event max_size in
-    let mk_part part part_bytes =
-      make_partial id (List.length parts) part part_bytes
-    in
-    parts |> List.mapi (fun i part_bytes -> marshal (mk_part i part_bytes))
+  let parts = split_bytes marshaled_point max_size in
+  let part_count = List.length parts in
+  let mk_part part part_bytes = { bytes = part_bytes; part; part_count } in
+  parts |> List.mapi (fun i part_bytes -> marshal (mk_part i part_bytes))
 
 (*****************************************************************************)
 (* Perf event *)
@@ -83,49 +74,70 @@ let perf_event_type =
 let perf_event =
   Runtime_events.User.register "Perf_event" Perf_event_tag perf_event_type
 
-let emit_event e =
-  let marshaled_events = marshal_event e in
+let emit_point (p : point) =
+  let marshaled_events = marshal_point p in
   List.iter
     (fun marshaled -> Runtime_events.User.write perf_event marshaled)
     marshaled_events
 [@@inline always]
 
-(* buffer for storing partial events so we can then rebuild them  *)
-type event_buffer = (Hash.t, (int * Bytes.t) list) Hashtbl.t
+(* buffer for storing partial points so we can then rebuild them  *)
+(* of type (ring_id, point parts) *)
+type point_buffer = (int, (int * Bytes.t) list) Hashtbl.t
+
+(** [event_of_perf_event ring_buffer_index buffer event] collects marshaled
+    events, and re-assembles them into points. Since the points are split into
+    parts, we return [None] if there was not enough parts to reconstruct a
+    point, or [Some point] if there were.
 
-(* event_of_perf event will keep track of partial events and spit them out as
-   normal events if it finds all the parts *)
-(* TODO this would probably make more sense if it spit out Point option or
-   similar? *)
-let event_of_perf_event buffer (marshaled, _) : t =
+    The runtime events file we read in has a unique ring buffer for each domain
+    ([ring_buffer_index]). Events are written to this buffer in order. This means
+    we can assume that the parts will come be read in order (e.g. ring 1 part 1,
+    ring 1 part 2 ...). We collect these parts to form the final point. If we
+    receive an out of order part, the last part in a point, or lose runtime events
+    in a ring buffer, we reset our point buffer *)
+let process_perf_event ring_buffer_index buffer (marshaled, _) : point option =
   let event = Marshal.from_bytes marshaled 0 in
+  let ring_parts =
+    match Hashtbl.find_opt buffer ring_buffer_index with
+    | Some parts -> parts
+    | None -> []
+  in
   match event with
-  | Partial { id; bytes; part_count; part } ->
-      let parts =
-        match Hashtbl.find_opt buffer id with
-        | Some parts -> (part, bytes) :: parts
-        | None -> [ (part, bytes) ]
-      in
-      let parts =
-        List.sort_uniq (fun (id1, _) (id2, _) -> Int.compare id1 id2) parts
-      in
-      if List.length parts = part_count then (
+  (* Don't store in buffer if we can immediately unmarshal *)
+  | { bytes; part_count; _ } when part_count = 1 ->
+      (* Also clear out the buffer just in case *)
+      Hashtbl.remove buffer ring_buffer_index;
+      Some (Marshal.from_bytes bytes 0)
+  (* If we don't have any parts, and receive something besides the start part,
+     just wait for the next start part*)
+  | { part; _ } when List.length ring_parts = 0 && part != 0 -> None
+  (* If we already have some parts, or this is the start part, begin collecting parts *)
+  | { bytes; part_count; _ } ->
+      let parts = bytes :: ring_parts in
+      let parts_len = List.length parts in
+      (* If we have enough then unmarshal! *)
+      (* TODO: We probably can just make the array all at once since we know the
+         size in theory? *)
+      if parts_len = part_count then (
         let full_bytes =
-          parts |> List.map snd
-          |> List.fold_left
-               (fun acc bytes ->
-                 let new_acc =
-                   Bytes.create (Bytes.length acc + Bytes.length bytes)
-                 in
-                 Bytes.blit acc 0 new_acc 0 (Bytes.length acc);
-                 Bytes.blit bytes 0 new_acc (Bytes.length acc)
-                   (Bytes.length bytes);
-                 new_acc)
-               (Bytes.create 0)
+          List.fold_right
+            (fun acc bytes ->
+              let new_acc =
+                Bytes.create (Bytes.length acc + Bytes.length bytes)
+              in
+              Bytes.blit acc 0 new_acc 0 (Bytes.length acc);
+              Bytes.blit bytes 0 new_acc (Bytes.length acc) (Bytes.length bytes);
+              new_acc)
+            parts (Bytes.create 0)
         in
-        Hashtbl.remove buffer id;
-        Marshal.from_bytes full_bytes 0)
+        Hashtbl.remove buffer ring_buffer_index;
+        Some (Marshal.from_bytes full_bytes 0))
+      else if parts_len < part_count then (
+        (* Weird state, clear buffer *)
+        Hashtbl.remove buffer ring_buffer_index;
+        None)
       else (
-        Hashtbl.replace buffer id parts;
-        event)
-  | _ -> event
+        (* If not then update the buffer *)
+        Hashtbl.replace buffer ring_buffer_index parts;
+        None)

lib/Pyro_caml_instruments.ml

@@ -15,10 +15,6 @@
 
 open Event
 
-let src = Logs.Src.create "pyro_caml" ~doc:"Pyro Caml"
-
-module Log = (val Logs.src_log src)
-
 (*****************************************************************************)
 (* Instrument side code *)
 (*****************************************************************************)
@@ -35,8 +31,8 @@ let emit_point_event raw_backtrace =
      that doesn't play well when linked into a rust program. Monotomic time
      would be nice so if the user/system changes the time of day we aren't
      screwed up, but for now we can assume that probably won't happen much*)
-  let event = Point (Unix.gettimeofday (), raw_stack_trace) in
-  emit_event event
+  let point = (Unix.gettimeofday (), raw_stack_trace) in
+  emit_point point
 [@@inline always]
 
 let tracker : (unit, unit) Gc.Memprof.tracker =
@@ -77,28 +73,33 @@ let maybe_with_memprof_sampler ?sampling_rate f =
 (* Profiler code *)
 (*****************************************************************************)
 let create_cursor path pid = Runtime_events.create_cursor (Some (path, pid))
-let empty_callbacks = Runtime_events.Callbacks.create ()
 
 (* Minimize work we do in process event since the instrumented program can write
    events quickly and so we need to keep pace while polling if we can *)
-let process_event now interval sample_points = function
-  | Point (time, raw_st) ->
+let process_point now interval sample_points = function
+  | Some (time, raw_st) ->
       if now -. time < interval then
         sample_points := (time, raw_st) :: !sample_points
-  | Partial _ -> ()
+  | None -> ()
 
-let read_poll ?(max_events = None) ?(callbacks = empty_callbacks) cursor
-    interval =
-  let event_buffer = Hashtbl.create 1000 in
+let read_poll ?(max_events = None) cursor interval =
+  let point_buffer = Hashtbl.create 1000 in
   let now = Unix.gettimeofday () in
   let sample_points = ref [] in
+  let callbacks =
+    Runtime_events.Callbacks.create
+      ~lost_events:(fun (ring_buffer_index : int) (_num_lost : int) ->
+        (* If we've lost events clear that ring buffer's event buffer *)
+        Hashtbl.remove point_buffer ring_buffer_index)
+      ()
+  in
   let callbacks =
     Runtime_events.Callbacks.add_user_event perf_event_type
-      (fun (_ring_buffer_index : int) (_ts : Runtime_events.Timestamp.t)
-           _event_t (e : marshaled) ->
+      (fun (ring_buffer_index : int) (_ts : Runtime_events.Timestamp.t) _event_t
+           (e : marshaled) ->
         e
-        |> event_of_perf_event event_buffer
-        |> process_event now interval sample_points)
+        |> process_perf_event ring_buffer_index point_buffer
+        |> process_point now interval sample_points)
       callbacks
   in
   (* TODO? Multithread this? *)

lib/Pyro_caml_instruments.mli

@@ -22,16 +22,10 @@ val create_cursor : string -> int -> Runtime_events.cursor
     from the given [path] and [pid]. *)
 
 val read_poll :
-  ?max_events:int option ->
-  ?callbacks:Runtime_events.Callbacks.t ->
-  Runtime_events.cursor ->
-  float ->
-  Stack_trace.t list
+  ?max_events:int option -> Runtime_events.cursor -> float -> Stack_trace.t list
 (** [read_poll cursor sample_interval] will read the profiling runtime events
     from the given cursor, and will give attempt to give a single
     {!Stack_trace.t} per every unique thread id, within [sample_interval] of the
     start time of this function call. If a sample is not within this interval,
     it will not be considered. Additionally it also ensures the samples chosen
-    are those closest to the start of the call of this function. Existing
-    callbacks can also be passed via [?callbacks], and those will be included
-    when reading the cursor *)
+    are those closest to the start of the call of this function. *)

lib/dune

@@ -1,7 +1,7 @@
 (library
  (name pyro_caml_instruments)
  (public_name pyro-caml-instruments)
- (libraries runtime_events unix logs digestif)
+ (libraries runtime_events unix logs)
  (preprocess
   (pps ppx_deriving.eq))
  (instrumentation.backend

pyro-caml-instruments.opam

@@ -11,7 +11,6 @@ depends: [
   "dune" {>= "3.20"}
   "ocaml"
   "logs"
-  "digestif"
   "ppx_deriving"
   "odoc" {with-doc}
 ]