Support running commands while attached

README.md

@@ -71,7 +71,7 @@ You can point magic-trace at a function such that when your application calls it
 
 1. [Here](https://raw.githubusercontent.com/janestreet/magic-trace/master/demo/demo.c)'s a sample C program to try out. It's a slightly modified version of the example in `man 3 dlopen`. Download that, build it with `gcc demo.c -ldl -o demo`, then leave it running `./demo`. We're going to use that program to learn how `dlopen` works.
 
-2. Run `magic-trace attach -pid $(pidof demo)`. When you see the message that it's successfully attached, wait a couple seconds and <kbd>Ctrl</kbd>+<kbd>C</kbd> `magic-trace`. It will output a file called `trace.fxt.gz` in your working directory.
+2. Run `magic-trace attach -pid $(pidof demo)`. When you see the message that it's successfully attached, wait a couple seconds and <kbd>Ctrl</kbd>+<kbd>C</kbd> `magic-trace`. It will output a file called `trace.fxt.gz` in your working directory. If you want to trace while running a specific non-interactive workload, pass the workload after `--`; for example, `magic-trace attach -pid $(pidof demo) -- ./load-generator`. The workload's standard output and standard error are forwarded, and magic-trace detaches when the workload exits. If recording stops first, magic-trace terminates the workload's process group.
 
 <p align="center">
   <img src="docs/assets/stage-1.gif">

src/attached_command/attached_command.ml

@@ -0,0 +1,127 @@
+open! Core
+open! Async
+
+type t =
+  { process : Process.t
+  ; finished : unit Or_error.t Deferred.t
+  }
+
+type outcome =
+  | Command_finished of unit Or_error.t
+  | Recording_stopped
+
+let create ~prog ~argv =
+  (* A separate process group lets us stop the workload and any descendants if the
+     recording ends first. [forward_output_and_wait] closes stdin and drains both output
+     streams, so this path is intentionally for non-interactive workloads. *)
+  let%map.Deferred process =
+    Process.create ~setpgid:Core_unix.Pgid.new_process_group ~prog ~args:argv ()
+  in
+  Or_error.map process ~f:(fun process ->
+    let finished = Process.forward_output_and_wait process in
+    { process; finished })
+;;
+
+let wait_for_outcome ~stop ~finished =
+  let%map.Deferred () =
+    Deferred.any_unit [ stop; Deferred.map finished ~f:(fun _ -> ()) ]
+  in
+  (* Prefer the command result if the command and recording finish in the same Async
+     cycle. In particular, this prevents a non-zero exit from being hidden by the
+     recording stop that the command itself initiates. *)
+  match Deferred.peek finished with
+  | Some result -> Command_finished result
+  | None -> Recording_stopped
+;;
+
+let wait t ~stop = wait_for_outcome ~stop ~finished:t.finished
+
+let send_to_process_group t signal =
+  Signal_unix.send_i signal (`Group (Process.pid t.process))
+;;
+
+let terminate ?(grace_period = Time_ns.Span.of_sec 1.) ?(signal = Signal.term) t =
+  send_to_process_group t signal;
+  match%bind.Deferred Clock_ns.with_timeout grace_period t.finished with
+  | `Result result -> return result
+  | `Timeout ->
+    send_to_process_group t Signal.kill;
+    t.finished
+;;
+
+module%test _ = struct
+  let%expect_test "command result wins a simultaneous stop" =
+    let stop = Ivar.create () in
+    let finished = Ivar.create () in
+    Ivar.fill_exn stop ();
+    Ivar.fill_exn finished (Ok ());
+    let%map outcome =
+      wait_for_outcome ~stop:(Ivar.read stop) ~finished:(Ivar.read finished)
+    in
+    (match outcome with
+     | Command_finished _ -> print_endline "command finished"
+     | Recording_stopped -> print_endline "recording stopped");
+    [%expect {| command finished |}]
+  ;;
+
+  let%expect_test "recording can stop before the command" =
+    let stop = Ivar.create () in
+    let finished = Ivar.create () in
+    Ivar.fill_exn stop ();
+    let%map outcome =
+      wait_for_outcome ~stop:(Ivar.read stop) ~finished:(Ivar.read finished)
+    in
+    (match outcome with
+     | Command_finished _ -> print_endline "command finished"
+     | Recording_stopped -> print_endline "recording stopped");
+    [%expect {| recording stopped |}]
+  ;;
+
+  let%expect_test "a non-zero command exit is preserved" =
+    let%bind command =
+      create
+        ~prog:"/bin/sh"
+        ~argv:
+          [ "-c"
+          ; "printf 'argument: %s\\n' \"$1\"; exit 23"
+          ; "magic-trace-test"
+          ; "with spaces"
+          ]
+    in
+    let command = ok_exn command in
+    let%map outcome = wait command ~stop:(Deferred.never ()) in
+    (match outcome with
+     | Recording_stopped -> print_endline "recording stopped"
+     | Command_finished result -> printf "command failed: %b\n" (Or_error.is_error result));
+    [%expect {|
+      argument: with spaces
+      command failed: true |}]
+  ;;
+
+  let%expect_test "stopping the recording reaps the command process group" =
+    let%bind command =
+      create ~prog:"/bin/sh" ~argv:[ "-c"; "trap '' TERM; sleep 60 & wait" ]
+    in
+    let command = ok_exn command in
+    (* Give the shell enough time to install its TERM handler and start the child. *)
+    let%bind () = Clock_ns.after (Time_ns.Span.of_ms 50.) in
+    let stop = Ivar.create () in
+    let outcome = wait command ~stop:(Ivar.read stop) in
+    Ivar.fill_exn stop ();
+    let%bind outcome in
+    (match outcome with
+     | Command_finished _ -> print_endline "command finished"
+     | Recording_stopped -> print_endline "recording stopped");
+    let%bind result = terminate command ~grace_period:(Time_ns.Span.of_ms 10.) in
+    let process_group_reaped =
+      match Signal_unix.send Signal.zero (`Group (Process.pid command.process)) with
+      | `No_such_process -> true
+      | `Ok -> false
+    in
+    printf "process group reaped: %b\n" (Or_error.is_error result && process_group_reaped);
+    [%expect {|
+      recording stopped
+      process group reaped: true |}]
+    |> Deferred.return
+  ;;
+end

src/attached_command/dune

@@ -0,0 +1,8 @@
+(library
+ (name magic_trace_attached_command)
+ (package magic-trace)
+ (wrapped false)
+ (libraries async core core_unix.signal_unix core_unix.sys_unix)
+ (inline_tests)
+ (preprocess
+  (pps ppx_jane)))

src/dune

@@ -42,6 +42,7 @@
   cohttp
   cohttp_static_handler
   core_unix.signal_unix
+  magic_trace_attached_command
   tracing
   magic_trace
   owee

src/trace.ml

@@ -538,6 +538,61 @@ module Make_commands (Backend : Backend_intf.S) = struct
     detach attachment
   ;;
 
+  let attach_and_record_while_running_command
+    record_opts
+    ~elf
+    ~debug_print_perf_commands
+    ~collection_mode
+    pids
+    ~prog
+    ~argv
+    =
+    let open Deferred.Or_error.Let_syntax in
+    let%bind attachment =
+      attach
+        record_opts
+        ~elf
+        ~debug_print_perf_commands
+        ~subcommand:Attach
+        ~collection_mode
+        pids
+    in
+    let { Attachment.done_ivar; _ } = attachment in
+    let stop = Ivar.read done_ivar in
+    let detach_with command_result =
+      let%map.Deferred detach_result = detach attachment in
+      Or_error.combine_errors_unit [ command_result; detach_result ]
+    in
+    let command_stop_signal = ref Signal.term in
+    Async_unix.Signal.handle ~stop [ Signal.int ] ~f:(fun (_ : Signal.t) ->
+      Core.eprintf "[ Got signal, detaching... ]\n%!";
+      command_stop_signal := Signal.int;
+      Ivar.fill_if_empty done_ivar ());
+    Deferred.upon stop (fun () -> Core.Signal.Expert.set Signal.int Default);
+    if Deferred.is_determined stop
+    then detach attachment
+    else (
+      Core.eprintf "[ Attached. Running command... ]\n%!";
+      match%bind.Deferred Attached_command.create ~prog ~argv with
+      | Error command_error ->
+        Ivar.fill_if_empty done_ivar ();
+        detach_with (Error command_error)
+      | Ok command ->
+        (match%bind.Deferred Attached_command.wait command ~stop with
+         | Attached_command.Command_finished command_result ->
+           Core.eprintf "[ Command exited, detaching... ]\n%!";
+           Ivar.fill_if_empty done_ivar ();
+           detach_with command_result
+         | Attached_command.Recording_stopped ->
+           Core.eprintf "[ Recording stopped, terminating command... ]\n%!";
+           let command_termination =
+             Attached_command.terminate command ~signal:!command_stop_signal
+           in
+           let%map.Deferred (_ : unit Or_error.t) = command_termination
+           and detach_result = detach attachment in
+           detach_result))
+  ;;
+
   let record_dir_flag mode =
     let open Command.Param in
     flag
@@ -605,6 +660,12 @@ module Make_commands (Backend : Backend_intf.S) = struct
     { Decode_opts.output_config; decode_opts; print_events }
   ;;
 
+  let command_argv_param ~escape_doc =
+    let%map_open.Command command = anon (maybe ("COMMAND" %: string))
+    and more_command = flag "--" escape ~doc:escape_doc in
+    Option.to_list command @ Option.value more_command ~default:[]
+  ;;
+
   let run_command =
     Command.async_or_error
       ~summary:"Runs a command and traces it."
@@ -622,11 +683,8 @@ module Make_commands (Backend : Backend_intf.S) = struct
        and decode_opts = decode_flags
        and debug_print_perf_commands
        and argv =
-         let%map_open.Command command = anon (maybe ("COMMAND" %: string))
-         and more_command =
-           flag "--" escape ~doc:"ARGS Arguments for the command. Ignored by magic-trace."
-         in
-         Option.to_list command @ Option.value more_command ~default:[]
+         command_argv_param
+           ~escape_doc:"ARGS Arguments for the command. Ignored by magic-trace."
        in
        fun () ->
          let open Deferred.Or_error.Let_syntax in
@@ -719,7 +777,9 @@ module Make_commands (Backend : Backend_intf.S) = struct
          magic-trace attach\n\n\
          # Fuzzy-find to select a running process and symbol to trigger on, snapshotting \
          the next time the symbol is called\n\
-         magic-trace attach -trigger ?\n")
+         magic-trace attach -trigger ?\n\n\
+         # Attach to a process, run a command, then detach when the command exits\n\
+         magic-trace attach -p $PID -- ./load-generator arg1 arg2\n")
       (let%map_open.Command record_opt_fn = record_flags
        and decode_opts = decode_flags
        and debug_print_perf_commands
@@ -731,6 +791,12 @@ module Make_commands (Backend : Backend_intf.S) = struct
            ~doc:
              "PID Processes to attach to as a comma separated list. Required if you \
               don't have the \"fzf\" application available in your PATH."
+       and command_argv =
+         command_argv_param
+           ~escape_doc:
+             "COMMAND Non-interactive command to run while attached. magic-trace \
+              detaches after the command exits and terminates it if recording stops \
+              first."
        in
        fun () ->
          let open Deferred.Or_error.Let_syntax in
@@ -759,12 +825,23 @@ module Make_commands (Backend : Backend_intf.S) = struct
                evaluate_trace_filter ~trace_filter:opts.trace_filter ~elf
              in
              let%bind () =
-               attach_and_record
-                 opts
-                 ~elf
-                 ~debug_print_perf_commands
-                 ~collection_mode
-                 pids
+               match command_argv with
+               | [] ->
+                 attach_and_record
+                   opts
+                   ~elf
+                   ~debug_print_perf_commands
+                   ~collection_mode
+                   pids
+               | prog :: argv ->
+                 attach_and_record_while_running_command
+                   opts
+                   ~elf
+                   ~debug_print_perf_commands
+                   ~collection_mode
+                   pids
+                   ~prog
+                   ~argv
              in
              let%bind.Deferred perf_maps = Perf_map.Table.load_by_pids pids in
              decode_to_trace