diff --git a/src/trace.ml b/src/trace.ml
index 4d0004703..38d4adb88 100644
--- a/src/trace.ml
+++ b/src/trace.ml
@@ -538,6 +538,46 @@ 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%bind.Deferred.Or_error 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
+ Async_unix.Signal.handle ~stop [ Signal.int ] ~f:(fun (_ : Signal.t) ->
+ Core.eprintf "[ Got signal, detaching... ]\n%!";
+ Ivar.fill_if_empty done_ivar ());
+ Deferred.upon stop (fun () -> Core.Signal.Expert.set Signal.int Default);
+ Core.eprintf "[ Attached. Running command... ]\n%!";
+ let command_finished = Process.run_forwarding ~prog ~args:argv () in
+ let command_result = Ivar.create () in
+ Deferred.upon command_finished (fun result ->
+ Core.eprintf "[ Command exited, detaching... ]\n%!";
+ Ivar.fill_if_empty done_ivar ();
+ Ivar.fill_if_empty command_result (Some result));
+ Deferred.upon stop (fun () -> Ivar.fill_if_empty command_result None);
+ let%bind.Deferred command_result = Ivar.read command_result in
+ let%bind () = stop in
+ let%bind () = detach attachment in
+ match command_result with
+ | None -> Deferred.Or_error.return ()
+ | Some result -> Deferred.return result
+ ;;
+
let record_dir_flag mode =
let open Command.Param in
flag
@@ -605,6 +645,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 +668,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 +762,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 +776,11 @@ 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 Command to run while attached. magic-trace detaches after the \
+ command exits."
in
fun () ->
let open Deferred.Or_error.Let_syntax in
@@ -759,12 +809,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
From f0191b696576cae700db6c999bcbfeca644d2b1a Mon Sep 17 00:00:00 2001
From: Kevin-Li-2025
diff --git a/src/attached_command/attached_command.ml b/src/attached_command/attached_command.ml
new file mode 100644
index 000000000..9ef0b3cfc
--- /dev/null
+++ b/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
diff --git a/src/attached_command/dune b/src/attached_command/dune
new file mode 100644
index 000000000..e13b681e5
--- /dev/null
+++ b/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)))
diff --git a/src/dune b/src/dune
index e9d36e436..f14998f94 100644
--- a/src/dune
+++ b/src/dune
@@ -42,6 +42,7 @@
cohttp
cohttp_static_handler
core_unix.signal_unix
+ magic_trace_attached_command
tracing
magic_trace
owee
diff --git a/src/trace.ml b/src/trace.ml
index 38d4adb88..c39c8ab2f 100644
--- a/src/trace.ml
+++ b/src/trace.ml
@@ -547,7 +547,8 @@ module Make_commands (Backend : Backend_intf.S) = struct
~prog
~argv
=
- let%bind.Deferred.Or_error attachment =
+ let open Deferred.Or_error.Let_syntax in
+ let%bind attachment =
attach
record_opts
~elf
@@ -558,24 +559,38 @@ module Make_commands (Backend : Backend_intf.S) = struct
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);
- Core.eprintf "[ Attached. Running command... ]\n%!";
- let command_finished = Process.run_forwarding ~prog ~args:argv () in
- let command_result = Ivar.create () in
- Deferred.upon command_finished (fun result ->
- Core.eprintf "[ Command exited, detaching... ]\n%!";
- Ivar.fill_if_empty done_ivar ();
- Ivar.fill_if_empty command_result (Some result));
- Deferred.upon stop (fun () -> Ivar.fill_if_empty command_result None);
- let%bind.Deferred command_result = Ivar.read command_result in
- let%bind () = stop in
- let%bind () = detach attachment in
- match command_result with
- | None -> Deferred.Or_error.return ()
- | Some result -> Deferred.return result
+ 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 =
@@ -779,8 +794,9 @@ module Make_commands (Backend : Backend_intf.S) = struct
and command_argv =
command_argv_param
~escape_doc:
- "COMMAND Command to run while attached. magic-trace detaches after the \
- command exits."
+ "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