Make gofmt look good

This is actually two separate hacks.

Hack 1:

gogo is either used to switch to something in the current stack or
an entirely separate stack.

If we're switching to an existing stack, assume that we're resuming
execution at the lowest stack frame with the same symbol as what gogo
jumped into.

If there's no existing stack frame that matches the gogo target,
assume we're switching to an entirely separate stack.

One place where this strategy breaks down is that it doesn't work
well with recursive functions that switch between goroutines. I
suspect that's not a very common pattern in non-adversaial go code,
and hope this will work well for most people.

Testing: There was already a test around this case, I'm trying out
this new strategy to cover more bases than just what's in the demo
script. Turns out real go code can jump into more than two
functions.

Hack 2:

A goroutine that returns, when there are no other goroutines,
returns into a function called "runtime.goexit.abi0".

When magic-trace sees a return into that symbol, it ends the
current thread and pretends there's a call to goexit. The goal
here is to end the outstanding stack frames and not infer goexit
as starting at the last thread end.

Testing: I added a perf test around the complete golang exit
sequence.

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. *)

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 |}]
+;;

test/golang_exit.mli

@@ -0,0 +1 @@
+(* intentionally left blank *)

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]

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