OxCaml Stack Switching

src/event.ml

@@ -88,6 +88,10 @@ module Ok = struct
           ; count : int
           ; name : Collection_mode.Event.Name.t
           }
+      | Ptwrite of
+          { location : Location.t
+          ; data : string
+          }
     [@@deriving sexp, bin_io]
   end
 

src/event.mli

@@ -56,6 +56,10 @@ module Ok : sig
           ; count : int
           ; name : Collection_mode.Event.Name.t
           } (** Represents counter based events collected through sampling. *)
+      | Ptwrite of
+          { location : Location.t
+          ; data : string
+          } (** Represents ptwrite events collected through Intel PT using PTWRITE. *)
     [@@deriving sexp]
   end
 

src/for_range.ml

@@ -25,7 +25,10 @@ let range_hit_times ~decode_events ~range_symbols =
     let is_start symbol = String.(Symbol.display_name symbol = start_symbol) in
     let is_stop symbol = String.(Symbol.display_name symbol = stop_symbol) in
     Pipe.filter_map events ~f:(function
-      | Error _ | Ok { data = Power _; _ } | Ok { data = Event_sample _; _ } -> None
+      | Error _
+      | Ok { data = Power _; _ }
+      | Ok { data = Event_sample _; _ }
+      | Ok { data = Ptwrite _; _ } -> None
       | Ok { data = Trace trace; time; _ } ->
         (match trace.kind with
          | Some Call ->
@@ -93,6 +96,7 @@ let decode_events_and_annotate ~decode_events ~range_symbols =
                | Ok { data = Power _; _ }
                | Ok { data = Stacktrace_sample _; _ }
                | Ok { data = Event_sample _; _ }
+               | Ok { data = Ptwrite _; _ }
                | Error _ -> hits, in_filtered_region)
           in
           ( (hits, in_filtered_region)

src/new_trace_writer.ml

@@ -59,16 +59,26 @@ module Thread_info = struct
     Trace_segment.add_event trace_segment event_data (Timestamp.create time)
   ;;
 
+  let set_fiber_id t fiber_id =
+    let #(trace_segment, ~in_filtered_region:_) = Nonempty_vec.last t.trace_segments in
+    Trace_segment.set_fiber_id trace_segment fiber_id
+  ;;
+
   module New_trace_segment_kind = struct
     type t =
       | Independent
       | Continuing_from_current
   end
 
-  let start_new_trace_segment t ~in_filtered_region ~(kind : New_trace_segment_kind.t) =
+  let start_new_trace_segment
+    t
+    ~in_filtered_region
+    ~(kind : New_trace_segment_kind.t)
+    ~fiber_stacks
+    =
     let new_trace_segment =
       match kind with
-      | Independent -> Trace_segment.create t.ocaml_exception_info
+      | Independent -> Trace_segment.create t.ocaml_exception_info fiber_stacks
       | Continuing_from_current ->
         let #(current, ~in_filtered_region:_) = Nonempty_vec.last t.trace_segments in
         Trace_segment.create_continuing_from current
@@ -89,6 +99,8 @@ type 'thread inner =
   ; annotate_inferred_start_times : bool
   ; mutable in_filtered_region : bool
   ; mutable transaction_events : Event.With_write_info.t Deque.t
+  ; pending_flows : ('thread -> Time_ns.Span.t -> unit) Hashtbl.M(Int).t
+  ; fiber_stacks : Trace_segment.Fiber.t Hashtbl.M(Int).t
   }
 
 type t = T : 'thread inner -> t
@@ -230,6 +242,8 @@ let create_expert
       ; annotate_inferred_start_times
       ; in_filtered_region = true
       ; transaction_events = Deque.create ()
+      ; pending_flows = Hashtbl.create (module Int)
+      ; fiber_stacks = Hashtbl.create (module Int)
       }
   in
   write_hits t hits;
@@ -336,7 +350,7 @@ let create_thread t event =
   ; extra_event_tracks = Hashtbl.create (module Collection_mode.Event.Name)
   ; trace_segments =
       Nonempty_vec.create
-        #( Trace_segment.create t.ocaml_exception_info
+        #( Trace_segment.create t.ocaml_exception_info t.fiber_stacks
          , ~in_filtered_region:t.in_filtered_region )
   }
 ;;
@@ -379,7 +393,8 @@ let maybe_start_filtered_region t ~should_write ~time:_ =
       Thread_info.start_new_trace_segment
         thread_info
         ~in_filtered_region:true
-        ~kind:Continuing_from_current);
+        ~kind:Continuing_from_current
+        ~fiber_stacks:t.fiber_stacks);
     t.in_filtered_region <- true)
 ;;
 
@@ -392,7 +407,8 @@ let maybe_stop_filtered_region t ~should_write =
       Thread_info.start_new_trace_segment
         thread_info
         ~in_filtered_region:false
-        ~kind:Continuing_from_current))
+        ~kind:Continuing_from_current
+        ~fiber_stacks:t.fiber_stacks))
 ;;
 
 let write_event_and_callstack (events_writer : Tracing_tool_output.events_writer) event =
@@ -501,6 +517,7 @@ and write_event' (T t) ?events_writer event =
       thread_info
       ~in_filtered_region:t.in_filtered_region
       ~kind:Independent
+      ~fiber_stacks:t.fiber_stacks
   | Ok event_value ->
     if should_write
     then
@@ -532,8 +549,8 @@ and write_event' (T t) ?events_writer event =
                  ~f:(fun pid -> [ "pid", Int (Pid.to_int pid) ])
                  ~default:[]
              ; Option.value_map
-                 (Event.thread outer_event).pid
-                 ~f:(fun pid -> [ "tid", Int (Pid.to_int pid) ])
+                 (Event.thread outer_event).tid
+                 ~f:(fun tid -> [ "tid", Int (Pid.to_int tid) ])
                  ~default:[]
              ])
        in
@@ -555,6 +572,52 @@ and write_event' (T t) ?events_writer event =
          ~name:"CPU"
          ~time
          ~args:Tracing.Trace.Arg.[ "freq (MHz)", Int freq ]
+     | { Event.Ok.thread = _ (* Already used this to look up thread info. *)
+       ; time = _
+       ; data = Ptwrite { location; data }
+       ; in_transaction = _
+       } ->
+       let symbol_name = Symbol.display_name location.symbol in
+       let is_perform = String.equal symbol_name "caml_perform" in
+       let is_resume = String.equal symbol_name "caml_resume" in
+       if is_perform || is_resume
+       then (
+         let fiber_id = Int64.to_int_trunc (Int64.of_string data) in
+         Thread_info.set_fiber_id thread_info fiber_id;
+         let name = if is_perform then "Perform Effect" else "Resume Continuation" in
+         let args = Tracing.Trace.Arg.[ "fiber", String (sprintf "0x%x" fiber_id) ] in
+         write_duration_complete t ~thread ~args ~name ~time ~time_end:time;
+         if is_perform
+         then (
+           let module T = (val t.trace) in
+           let flow = T.create_flow () in
+           T.write_flow_step flow ~thread ~time:(time :> Time_ns.Span.t);
+           Hashtbl.set t.pending_flows ~key:fiber_id ~data:(fun thread time ->
+             T.write_flow_step flow ~thread ~time;
+             T.finish_flow flow))
+         else (
+           match Hashtbl.find_and_remove t.pending_flows fiber_id with
+           | Some finish -> finish thread (time :> Time_ns.Span.t)
+           | None -> ()))
+       else (
+         let args =
+           Tracing.Trace.Arg.(
+             List.concat
+               [ [ "timestamp", Int (Time_ns.Span.to_int_ns (time :> Time_ns.Span.t)) ]
+               ; [ "symbol", String symbol_name ]
+               ; [ "addr", Pointer location.instruction_pointer ]
+               ; [ "data", String data ]
+               ; Option.value_map
+                   (Event.thread outer_event).pid
+                   ~f:(fun pid -> [ "pid", Int (Pid.to_int pid) ])
+                   ~default:[]
+               ; Option.value_map
+                   (Event.thread outer_event).tid
+                   ~f:(fun tid -> [ "tid", Int (Pid.to_int tid) ])
+                   ~default:[]
+               ])
+         in
+         write_duration_complete t ~thread ~args ~name:"PTWRITE" ~time ~time_end:time)
      | { Event.Ok.data = Stacktrace_sample _; _ } ->
        (* This should be unreachable, we currently delegate support for sampling to the old trace-writer. *)
        assert false)

src/perf_capabilities.ml

@@ -9,6 +9,7 @@ let snapshot_on_exit = bit 3
 let last_branch_record = bit 4
 let dlfilter = bit 5
 let ctlfd = bit 6
+let ptwrite = bit 7
 
 include Flags.Make (struct
     let allow_intersecting = false
@@ -22,6 +23,7 @@ include Flags.Make (struct
       ; last_branch_record, "last_branch_record"
       ; dlfilter, "dlfilter"
       ; ctlfd, "ctlfd"
+      ; ptwrite, "ptwrite"
       ]
     ;;
   end)
@@ -58,6 +60,18 @@ let supports_configurable_psb_period () =
   | Sys_error _ -> false
 ;;
 
+let supports_ptwrite () =
+  try
+    let ptwrite_cap =
+      In_channel.read_all "/sys/bus/event_source/devices/intel_pt/caps/ptwrite"
+    in
+    String.( = ) ptwrite_cap "1\n"
+  with
+  (* Even if this file is not present (i.e. when Intel PT isn't present), we
+     don't want capability checking to fail. *)
+  | Sys_error _ -> false
+;;
+
 (* This checks if pdcm flag is present in /proc/cpuinfo. This is necessary for
    LBR to work. Although I couldn't ascertain that it is also sufficient.
    However it seems unlikely this would fail on most machines. *)
@@ -113,6 +127,7 @@ let detect_exn () =
   let set_if bool flag cap = cap + if bool then flag else empty in
   empty
   |> set_if (supports_configurable_psb_period ()) configurable_psb_period
+  |> set_if (supports_ptwrite ()) ptwrite
   |> set_if (supports_tracing_kernel ()) kernel_tracing
   |> set_if (supports_kcore version) kcore
   |> set_if (supports_snapshot_on_exit version) snapshot_on_exit

src/perf_capabilities.mli

@@ -12,4 +12,5 @@ val snapshot_on_exit : t
 val last_branch_record : t
 val dlfilter : t
 val ctlfd : t
+val ptwrite : t
 val detect_exn : unit -> t Deferred.t

src/perf_decode.ml

@@ -29,6 +29,12 @@ let perf_cbr_event_re =
   Re.Perl.re {|^ *([a-z )(]*)? +cbr: +([0-9]+ +freq: +([0-9]+) MHz)?(.*)$|} |> Re.compile
 ;;
 
+let perf_ptwrite_event_re =
+  Re.Perl.re
+    {|^ *(call|return|tr strt|syscall|sysret|hw int|iret|int|tx abrt|tr end|tr strt tr end|tr end  (?:async|call|return|syscall|sysret|iret)|jmp|jcc)? +IP: +([0-9a-f]+) +payload: +((?:0x)?[0-9a-f]+) (.*) ([0-9]+) +([0-9a-f]+) (.*)$|}
+  |> Re.compile
+;;
+
 let trace_error_re =
   Re.Posix.re
     {|^ instruction trace error type [0-9]+ (time ([0-9]+)\.([0-9]+) )?cpu [\-0-9]+ pid ([\-0-9]+) tid ([\-0-9]+) ip (0x[0-9a-fA-F]+|0) code [0-9]+: (.*)$|}
@@ -47,7 +53,15 @@ type header =
       { thread : Event.Thread.t
       ; time : Time_ns.Span.t
       ; period : int
-      ; event : [ `Branches | `Cbr | `Psb | `Cycles | `Branch_misses | `Cache_misses ]
+      ; event :
+          [ `Branches
+          | `Cbr
+          | `Psb
+          | `Cycles
+          | `Branch_misses
+          | `Cache_misses
+          | `Ptwrite
+          ]
       ; remaining_line : string
       }
 
@@ -99,6 +113,7 @@ let parse_event_header line =
         | "cycles" -> `Cycles
         | "branch-misses" -> `Branch_misses
         | "cache-misses" -> `Cache_misses
+        | "ptwrite" -> `Ptwrite
         | _ ->
           raise_s
             [%message
@@ -349,6 +364,33 @@ let parse_perf_extra_sampled_event
     }
 ;;
 
+let parse_perf_ptwrite_event ?perf_maps (thread : Event.Thread.t) time line : Event.t =
+  match Re.Group.all (Re.exec perf_ptwrite_event_re line) with
+  | [| _
+     ; _branch_kind
+     ; _fup_ip
+     ; payload
+     ; _payload_ascii
+     ; _unknown
+     ; instruction_pointer
+     ; symbol_and_offset
+    |] ->
+    let location =
+      parse_location ?perf_maps ~pid:thread.pid instruction_pointer symbol_and_offset
+    in
+    Ok
+      { thread
+      ; time
+      ; data = Ptwrite { location; data = payload }
+      ; in_transaction = false
+      }
+  | results ->
+    raise_s
+      [%message
+        "Regex of perf ptwrite event did not match expected fields"
+          (results : string array)]
+;;
+
 let to_event ?perf_maps lines : Event.t option =
   try
     match lines with
@@ -385,7 +427,9 @@ let to_event ?perf_maps lines : Event.t option =
                  period
                  remaining_line
                  lines
-                 Cache_misses)))
+                 Cache_misses)
+          | `Ptwrite ->
+            Some (parse_perf_ptwrite_event ?perf_maps thread time remaining_line)))
   with
   | exn ->
     raise_s
@@ -755,6 +799,58 @@ module%test _ = struct
             (Trace (trace_state_change End) (kind Async) (src 0x7f6fce0b71f4)
              (dst 0x0)))))) |}]
   ;;
+
+  let%expect_test "perf ptwrite event" =
+    check
+      {| 2769074/2769293 2592496.569782106:          1                                        ptwrite:   call                   IP: 0 payload: 0x1b5                            0     55c7952ad2d0 [unknown] (foo.bin)|};
+    [%expect
+      {|
+        ((Ok
+          ((thread ((pid (2769074)) (tid (2769293)))) (time 30d8m16.569782106s)
+           (data (Ptwrite (location 0x55c7952ad2d0) (data 0x1b5)))))) |}]
+  ;;
+
+  let%expect_test "perf ptwrite event with printable payload" =
+    check
+      {|2817453/2817483 2596547.813541321:          1                                        ptwrite:   call                   IP: 0 payload: 0x62 b                           0     613d946b42d0 [unknown] (foo.bin)|};
+    [%expect
+      {|
+        ((Ok
+          ((thread ((pid (2817453)) (tid (2817483)))) (time 30d1h15m47.813541321s)
+           (data (Ptwrite (location 0x613d946b42d0) (data 0x62)))))) |}]
+  ;;
+
+  let%expect_test "perf ptwrite event with jcc instruction" =
+    check
+      {|3256341/3256760 2632746.069047397:          1                                        ptwrite:   jcc                    IP: 0 payload: 0x5000000000061b2                 0     577bcad80555 [unknown] (foo.bin)|};
+    [%expect
+      {|
+      ((Ok
+        ((thread ((pid (3256341)) (tid (3256760)))) (time 30d11h19m6.069047397s)
+         (data (Ptwrite (location 0x577bcad80555) (data 0x5000000000061b2)))))) |}]
+  ;;
+
+  let%expect_test "perf ptwrite event without instruction" =
+    check
+      {|3285832/3286108 2634666.172476558:          1                                        ptwrite:                          IP: 0 payload: 0xa00000000000000                 0     64d20fb10432 [unknown] (foo.bin)|};
+    [%expect
+      {|
+      ((Ok
+        ((thread ((pid (3285832)) (tid (3286108)))) (time 30d11h51m6.172476558s)
+         (data (Ptwrite (location 0x64d20fb10432) (data 0xa00000000000000))))))
+      |}]
+  ;;
+
+  let%expect_test "perf ptwrite event zero payload" =
+    check
+      {|997136/997136  340026.788189283:          1                                        ptwrite:   jcc                    IP: 0 payload: 0                 0     64d20fb10432 [unknown] (foo.bin)|};
+    [%expect
+      {|
+      ((Ok
+        ((thread ((pid (997136)) (tid (997136)))) (time 3d22h27m6.788189283s)
+         (data (Ptwrite (location 0x64d20fb10432) (data 0))))))
+      |}]
+  ;;
 end
 
 module For_testing = struct