janestreet · cgaebel · Apr 26, 2022
diff --git a/src/trace_writer.ml b/src/trace_writer.ml
@@ -1,6 +1,12 @@
 open! Core
 open! Import
 
+(* Whether a language-specific hack override the default behavior or not. *)
+type hack_result =
+  | Handled
+  | Ignored
+[@@deriving sexp_of]
+
 let debug = ref false
 let is_kernel_address addr = Int64.(addr < 0L)
 
@@ -479,7 +485,7 @@ let call t thread_info ~time ~location =
   Callstack.push thread_info.callstack location.symbol
 ;;
 
-let ret_without_checking_for_go_hacks t (thread_info : _ Thread_info.t) ~time =
+let ret t (thread_info : _ Thread_info.t) ~time =
   match Callstack.pop thread_info.callstack with
   | Some symbol -> add_event t thread_info time { symbol; kind = Ret }
   | None ->
@@ -495,7 +501,6 @@ let ret_without_checking_for_go_hacks t (thread_info : _ Thread_info.t) ~time =
 ;;
 
 let rec clear_callstack t (thread_info : _ Thread_info.t) ~time =
-  let ret = ret_without_checking_for_go_hacks in
   match Callstack.top thread_info.callstack with
   | None -> ()
   | Some _ ->
@@ -534,57 +539,79 @@ let end_of_thread t (thread_info : _ Thread_info.t) ~time ~is_kernel_address : u
 
    At startup (and maybe other situations?), gogo clears all callstacks and executes [main]. *)
 module Go_hacks : sig
-  val ret_track_gogo
+  val detect_stack_switch
+    :  'a inner
+    -> 'a Thread_info.t
+    -> time:Mapped_time.t
+    -> top_of_stack:Symbol.t
+    -> jump_target:Event.Location.t
+    -> hack_result
+
+  val detect_goexit
     :  'a inner
     -> 'a Thread_info.t
     -> time:Mapped_time.t
-    -> returned_from:Symbol.t option
-    -> unit
+    -> jump_target:Event.Location.t
+    -> hack_result
 end = struct
   let is_gogo (symbol : Symbol.t) =
     match symbol with
     | From_perf "gogo" -> true
     | _ -> false
   ;;
 
-  let is_known_gogo_destination (symbol : Symbol.t) =
+  let is_goexit_abi0 (symbol : Symbol.t) =
     match symbol with
-    | From_perf ("runtime.mcall" | "runtime.morestack.abi0") -> true
+    | From_perf "runtime.goexit.abi0" -> true
     | _ -> false
   ;;
 
-  let current_stack_contains_known_gogo_destination (thread_info : _ Thread_info.t) =
-    Stack.find thread_info.callstack.stack ~f:is_known_gogo_destination |> Option.is_some
+  (* Walk up the stack, popping stack frames until we hit gogo's jump target. If we make
+     it to the end of the stack without seeing a gogo jump target, assume that we're
+     switching to a separate stack. *)
+  let rec pop_until_gogo_destination t (thread_info : _ Thread_info.t) ~time ~jump_target =
+    if Callstack.is_empty thread_info.callstack
+    then (
+      (* Popped off the complete stack. If we hit this, there was a gogo but we didn't find
+       the destination function anywhere in the stack. We assume this means we're switching
+       to another stack entirely, so mark it as the end of the thread so future decode errors
+       get marked to this point and past decode errors are flushed. *)
+      end_of_thread t thread_info ~time ~is_kernel_address:false;
+      `switched_stacks)
+    else (
+      ret t thread_info ~time;
+      if [%compare.equal: Symbol.t option]
+           (Callstack.top thread_info.callstack)
+           (Some jump_target)
+      then `jumped_internally
+      else pop_until_gogo_destination t thread_info ~time ~jump_target)
   ;;
 
-  let rec pop_until_gogo_destination t (thread_info : _ Thread_info.t) ~time =
-    let ret = ret_without_checking_for_go_hacks in
-    match Callstack.top thread_info.callstack with
-    | None -> ()
-    | Some symbol ->
-      ret t thread_info ~time;
-      (* Return one past the known gogo destination. This hack is necessary because:
-
-         - all gogo-destination functions are jumped into and out of
-         - magic-trace translates the jump returning from gogo-destination into a ret/call pair
-         - this runs on the ret, but the call is to gogo-destination's caller and we don't
-           want a second stack frame for that.
-
-         This is a little janky because you see a stack frame momentarily end then start back
-         up again on every [gogo]. I think that's a small price to pay to keep all the Go hacks
-         in one place. *)
-      if is_known_gogo_destination symbol
-      then ret t thread_info ~time
-      else pop_until_gogo_destination t thread_info ~time
+  let detect_stack_switch
+      t
+      thread_info
+      ~time
+      ~top_of_stack
+      ~(jump_target : Event.Location.t)
+    =
+    if is_gogo top_of_stack
+    then (
+      (match
+         pop_until_gogo_destination t thread_info ~time ~jump_target:jump_target.symbol
+       with
+      | `jumped_internally -> ()
+      | `switched_stacks -> call t thread_info ~time ~location:jump_target);
+      Handled)
+    else Ignored
   ;;
 
-  let ret_track_gogo t thread_info ~time ~returned_from =
-    let is_ret_from_gogo = Option.value_map ~f:is_gogo returned_from ~default:false in
-    if is_ret_from_gogo
-    then
-      if current_stack_contains_known_gogo_destination thread_info
-      then pop_until_gogo_destination t thread_info ~time
-      else end_of_thread t thread_info ~time ~is_kernel_address:false
+  let detect_goexit t thread_info ~time ~(jump_target : Event.Location.t) =
+    if is_goexit_abi0 jump_target.symbol
+    then (
+      end_of_thread t thread_info ~time ~is_kernel_address:false;
+      call t thread_info ~time ~location:jump_target;
+      Handled)
+    else Ignored
   ;;
 end
 
@@ -596,9 +623,6 @@ end
 module Ocaml_hacks : sig
   val ret_track_exn_data : 'a inner -> 'a Thread_info.t -> time:Mapped_time.t -> unit
 end = struct
-  (* It's ocaml, not go. *)
-  let ret = ret_without_checking_for_go_hacks
-
   let unwind_stack t (thread_info : _ Thread_info.t) ~time diff =
     let frames_to_unwind = thread_info.frames_to_unwind in
     for _ = 0 to !frames_to_unwind + diff do
@@ -621,12 +645,6 @@ end = struct
   ;;
 end
 
-let ret t (thread_info : _ Thread_info.t) ~time : unit =
-  let returned_from = Callstack.top thread_info.callstack in
-  ret_without_checking_for_go_hacks t thread_info ~time;
-  Go_hacks.ret_track_gogo t thread_info ~time ~returned_from
-;;
-
 let check_current_symbol
     t
     (thread_info : _ Thread_info.t)
@@ -638,10 +656,21 @@ let check_current_symbol
      with jumps between functions (e.g. tailcalls, PLT) or returns out of the highest
      known function, so we have to correct the top of the stack here. *)
   match Callstack.top thread_info.callstack with
-  | Some known when not ([%compare.equal: Symbol.t] known location.symbol) ->
-    ret t thread_info ~time;
-    call t thread_info ~time ~location
-  | Some _ -> ()
+  | Some known ->
+    if not ([%compare.equal: Symbol.t] known location.symbol)
+    then (
+      match
+        Go_hacks.detect_stack_switch
+          t
+          thread_info
+          ~time
+          ~top_of_stack:known
+          ~jump_target:location
+      with
+      | Handled -> ()
+      | Ignored ->
+        ret t thread_info ~time;
+        call t thread_info ~time ~location)
   | None ->
     (* If we have no callstack left, then we just returned out of something we didn't
        see the call for. Since we're in snapshot mode, this happens with functions
@@ -729,8 +758,11 @@ let write_event (T t) event =
       call t thread_info ~time ~location:Event.Location.syscall
     | Some Return, Some End -> call t thread_info ~time ~location:Event.Location.returned
     | Some Return, None ->
-      Ocaml_hacks.ret_track_exn_data t thread_info ~time;
-      check_current_symbol t thread_info ~time dst
+      (match Go_hacks.detect_goexit t thread_info ~time ~jump_target:dst with
+      | Handled -> ()
+      | Ignored ->
+        Ocaml_hacks.ret_track_exn_data t thread_info ~time;
+        check_current_symbol t thread_info ~time dst)
     | None, Some Start ->
       (* Might get this under /u, /k, and /uk, but we need to handle them all
        differently. *)

diff --git a/test/golang_exit.ml b/test/golang_exit.ml
@@ -0,0 +1,67 @@
+open! Core
+
+let%expect_test "golang exit sequence" =
+  Perf_script.run ~trace_mode:Userspace "golang_exit.perf";
+  [%expect
+    {|
+    3268985/3268994 3444853.363323631:   call                           439440 runtime.execute+0x80 =>           4365e0 runtime.casgstatus+0x0
+    3268985/3268994 3444853.363323637:   return                         4368ab runtime.casgstatus+0x2cb =>           439445 runtime.execute+0x85
+    ->      1ns BEGIN runtime.casgstatus
+    3268985/3268994 3444853.363323641:   call                           4394e9 runtime.execute+0x129 =>           45e100 runtime.gogo.abi0+0x0
+    ->      7ns END   runtime.casgstatus
+    3268985/3268994 3444853.363323642:   jmp                            45e10c runtime.gogo.abi0+0xc =>           45d320 gogo+0x0
+    ->     11ns BEGIN runtime.gogo.abi0
+    3268985/3268994 3444853.363323643:   jmp                            45d35e gogo+0x3e =>           434b36 runtime.gopark+0xd6
+    ->     12ns END   runtime.gogo.abi0
+    ->     12ns BEGIN gogo
+    ->      0ns BEGIN runtime.execute [inferred start time]
+    ->     13ns END   gogo
+    ->     13ns END   runtime.execute
+    3268985/3268994 3444853.363323644:   return                         434b40 runtime.gopark+0xe0 =>           4061cc runtime.chanrecv+0x56c
+    ->     13ns BEGIN runtime.gopark
+    3268985/3268994 3444853.363323648:   call                           406280 runtime.chanrecv+0x620 =>           434f60 runtime.releaseSudog+0x0
+    ->     14ns END   runtime.gopark
+    3268985/3268994 3444853.363323652:   return                         4350ce runtime.releaseSudog+0x16e =>           406285 runtime.chanrecv+0x625
+    ->     18ns BEGIN runtime.releaseSudog
+    3268985/3268994 3444853.363323653:   return                         40629e runtime.chanrecv+0x63e =>           405c38 runtime.chanrecv1+0x18
+    ->     22ns END   runtime.releaseSudog
+    3268985/3268994 3444853.363323654:   return                         405c41 runtime.chanrecv1+0x21 =>           507a72 main.(*sequencer).GetExitCode+0x92
+    ->     23ns END   runtime.chanrecv
+    3268985/3268994 3444853.363323655:   return                         507a80 main.(*sequencer).GetExitCode+0xa0 =>           508f0f main.main+0x14f
+    ->     24ns END   runtime.chanrecv1
+    3268985/3268994 3444853.363323656:   call                           508f0f main.main+0x14f =>           49d300 os.Exit+0x0
+    ->     25ns END   main.(*sequencer).GetExitCode
+    3268985/3268994 3444853.363323657:   call                           49d320 os.Exit+0x20 =>           495f40 internal/testlog.PanicOnExit0+0x0
+    ->     26ns BEGIN os.Exit
+    3268985/3268994 3444853.363323659:   call                           495fd9 internal/testlog.PanicOnExit0+0x99 =>           496020 internal/testlog.PanicOnExit0.func1+0x0
+    ->     27ns BEGIN internal/testlog.PanicOnExit0
+    3268985/3268994 3444853.363323660:   call                           496041 internal/testlog.PanicOnExit0.func1+0x21 =>           465d20 sync.(*Mutex).Unlock+0x0
+    ->     29ns BEGIN internal/testlog.PanicOnExit0.func1
+    3268985/3268994 3444853.363323662:   return                         465d52 sync.(*Mutex).Unlock+0x32 =>           496046 internal/testlog.PanicOnExit0.func1+0x26
+    ->     30ns BEGIN sync.(*Mutex).Unlock
+    3268985/3268994 3444853.363323663:   return                         49604f internal/testlog.PanicOnExit0.func1+0x2f =>           495fdb internal/testlog.PanicOnExit0+0x9b
+    ->     32ns END   sync.(*Mutex).Unlock
+    3268985/3268994 3444853.363323664:   return                         495fe9 internal/testlog.PanicOnExit0+0xa9 =>           49d325 os.Exit+0x25
+    ->     33ns END   internal/testlog.PanicOnExit0.func1
+    3268985/3268994 3444853.363323665:   call                           49d329 os.Exit+0x29 =>           45bc20 os.runtime_beforeExit+0x0
+    ->     34ns END   internal/testlog.PanicOnExit0
+    3268985/3268994 3444853.363323666:   return                         45bc20 os.runtime_beforeExit+0x0 =>           49d32e os.Exit+0x2e
+    ->     35ns BEGIN os.runtime_beforeExit
+    3268985/3268994 3444853.363323667:   call                           49d333 os.Exit+0x33 =>           45c300 syscall.Exit+0x0
+    ->     36ns END   os.runtime_beforeExit
+    3268985/3268994 3444853.363323668:   call                           45c311 syscall.Exit+0x11 =>           4619e0 runtime.exit.abi0+0x0
+    ->     37ns BEGIN syscall.Exit
+    3268985/3268994 3444853.363323669:   tr end  syscall                4619e9 runtime.exit.abi0+0x9 =>                0 [unknown]
+    ->     38ns BEGIN runtime.exit.abi0
+    END
+    ->     13ns BEGIN main.main [inferred start time]
+    ->     13ns BEGIN main.(*sequencer).GetExitCode [inferred start time]
+    ->     13ns BEGIN runtime.chanrecv1 [inferred start time]
+    ->     13ns BEGIN runtime.chanrecv [inferred start time]
+    ->     39ns BEGIN [syscall]
+    ->     39ns END   [syscall]
+    ->     39ns END   runtime.exit.abi0
+    ->     39ns END   syscall.Exit
+    ->     39ns END   os.Exit
+    ->     39ns END   main.main |}]
+;;
diff --git a/test/golang_exit.mli b/test/golang_exit.mli
@@ -0,0 +1 @@
+(* intentionally left blank *)
diff --git a/test/golang_exit.perf b/test/golang_exit.perf
@@ -0,0 +1,41 @@
+3268985/3268994 3444853.363323630:   jmp                            439421 runtime.execute+0x61 =>           43942c runtime.execute+0x6c
+3268985/3268994 3444853.363323631:   call                           439440 runtime.execute+0x80 =>           4365e0 runtime.casgstatus+0x0
+3268985/3268994 3444853.363323632:   jcc                            43660a runtime.casgstatus+0x2a =>           43664e runtime.casgstatus+0x6e
+3268985/3268994 3444853.363323633:   jmp                            436652 runtime.casgstatus+0x72 =>           43666a runtime.casgstatus+0x8a
+3268985/3268994 3444853.363323634:   jcc                            436683 runtime.casgstatus+0xa3 =>           43672a runtime.casgstatus+0x14a
+3268985/3268994 3444853.363323635:   jcc                            43672d runtime.casgstatus+0x14d =>           436751 runtime.casgstatus+0x171
+3268985/3268994 3444853.363323636:   jcc                            436758 runtime.casgstatus+0x178 =>           4368a2 runtime.casgstatus+0x2c2
+3268985/3268994 3444853.363323637:   return                         4368ab runtime.casgstatus+0x2cb =>           439445 runtime.execute+0x85
+3268985/3268994 3444853.363323638:   jmp                            439479 runtime.execute+0xb9 =>           43948e runtime.execute+0xce
+3268985/3268994 3444853.363323639:   jcc                            4394a6 runtime.execute+0xe6 =>           4394b2 runtime.execute+0xf2
+3268985/3268994 3444853.363323640:   jcc                            4394b9 runtime.execute+0xf9 =>           4394e1 runtime.execute+0x121
+3268985/3268994 3444853.363323641:   call                           4394e9 runtime.execute+0x129 =>           45e100 runtime.gogo.abi0+0x0
+3268985/3268994 3444853.363323642:   jmp                            45e10c runtime.gogo.abi0+0xc =>           45d320 gogo+0x0
+3268985/3268994 3444853.363323643:   jmp                            45d35e gogo+0x3e =>           434b36 runtime.gopark+0xd6
+3268985/3268994 3444853.363323644:   return                         434b40 runtime.gopark+0xe0 =>           4061cc runtime.chanrecv+0x56c
+3268985/3268994 3444853.363323645:   jmp                            406201 runtime.chanrecv+0x5a1 =>           40620f runtime.chanrecv+0x5af
+3268985/3268994 3444853.363323646:   jcc                            40621d runtime.chanrecv+0x5bd =>           40623e runtime.chanrecv+0x5de
+3268985/3268994 3444853.363323647:   jmp                            406262 runtime.chanrecv+0x602 =>           40627c runtime.chanrecv+0x61c
+3268985/3268994 3444853.363323648:   call                           406280 runtime.chanrecv+0x620 =>           434f60 runtime.releaseSudog+0x0
+3268985/3268994 3444853.363323648:   jcc                            435007 runtime.releaseSudog+0xa7 =>           435012 runtime.releaseSudog+0xb2
+3268985/3268994 3444853.363323649:   jcc                            43502e runtime.releaseSudog+0xce =>           435079 runtime.releaseSudog+0x119
+3268985/3268994 3444853.363323650:   jmp                            435095 runtime.releaseSudog+0x135 =>           43509c runtime.releaseSudog+0x13c
+3268985/3268994 3444853.363323651:   jcc                            4350bb runtime.releaseSudog+0x15b =>           4350c5 runtime.releaseSudog+0x165
+3268985/3268994 3444853.363323652:   return                         4350ce runtime.releaseSudog+0x16e =>           406285 runtime.chanrecv+0x625
+3268985/3268994 3444853.363323653:   return                         40629e runtime.chanrecv+0x63e =>           405c38 runtime.chanrecv1+0x18
+3268985/3268994 3444853.363323654:   return                         405c41 runtime.chanrecv1+0x21 =>           507a72 main.(*sequencer).GetExitCode+0x92
+3268985/3268994 3444853.363323655:   return                         507a80 main.(*sequencer).GetExitCode+0xa0 =>           508f0f main.main+0x14f
+3268985/3268994 3444853.363323656:   call                           508f0f main.main+0x14f =>           49d300 os.Exit+0x0
+3268985/3268994 3444853.363323657:   call                           49d320 os.Exit+0x20 =>           495f40 internal/testlog.PanicOnExit0+0x0
+3268985/3268994 3444853.363323658:   jcc                            495f86 internal/testlog.PanicOnExit0+0x46 =>           495f94 internal/testlog.PanicOnExit0+0x54
+3268985/3268994 3444853.363323659:   call                           495fd9 internal/testlog.PanicOnExit0+0x99 =>           496020 internal/testlog.PanicOnExit0.func1+0x0
+3268985/3268994 3444853.363323660:   call                           496041 internal/testlog.PanicOnExit0.func1+0x21 =>           465d20 sync.(*Mutex).Unlock+0x0
+3268985/3268994 3444853.363323661:   jcc                            465d42 sync.(*Mutex).Unlock+0x22 =>           465d49 sync.(*Mutex).Unlock+0x29
+3268985/3268994 3444853.363323662:   return                         465d52 sync.(*Mutex).Unlock+0x32 =>           496046 internal/testlog.PanicOnExit0.func1+0x26
+3268985/3268994 3444853.363323663:   return                         49604f internal/testlog.PanicOnExit0.func1+0x2f =>           495fdb internal/testlog.PanicOnExit0+0x9b
+3268985/3268994 3444853.363323664:   return                         495fe9 internal/testlog.PanicOnExit0+0xa9 =>           49d325 os.Exit+0x25
+3268985/3268994 3444853.363323665:   call                           49d329 os.Exit+0x29 =>           45bc20 os.runtime_beforeExit+0x0
+3268985/3268994 3444853.363323666:   return                         45bc20 os.runtime_beforeExit+0x0 =>           49d32e os.Exit+0x2e
+3268985/3268994 3444853.363323667:   call                           49d333 os.Exit+0x33 =>           45c300 syscall.Exit+0x0
+3268985/3268994 3444853.363323668:   call                           45c311 syscall.Exit+0x11 =>           4619e0 runtime.exit.abi0+0x0
+3268985/3268994 3444853.363323669:   tr end  syscall                4619e9 runtime.exit.abi0+0x9 =>                0 [unknown]
diff --git a/test/simple_gogo.ml b/test/simple_gogo.ml
@@ -43,8 +43,6 @@ let%expect_test "an mcall/gogo pair from a small go program" =
     ->  70.97us END   runtime.schedule
     ->  70.97us END   runtime.park_m
     ->  70.97us END   runtime.mcall
-    ->  70.97us END   runtime.gopark
-    ->  70.97us BEGIN runtime.gopark
     3109264/3109264 3363099.850069207:   call                           404ea0 runtime.chanrecv+0x620 =>           430ce0 runtime.releaseSudog+0x0
     -> 70.982us END   runtime.gopark
     3109264/3109264 3363099.850069390:   return                         430e4e runtime.releaseSudog+0x16e =>           404ea5 runtime.chanrecv+0x625