Stop rbspy and its command when singed is terminated

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    singed (0.3.0)
+    singed (0.4.0)
       stackprof (>= 0.2.13)
 
 GEM

README.md

@@ -196,6 +196,10 @@ $ bundle exec singed -- bin/rails runner 'Model.all.to_a'
 
 The flamegraph is opened afterwards.
 
+To profile a command that runs until it's stopped, like a server, stop it with Ctrl-C. Or, when `singed` runs in the background, such as from a script, stop it with `kill`'s default SIGTERM. Either way, rbspy stops the command and writes the flamegraph, which `singed` then opens. `singed` ignores a SIGINT sent to it alone, such as by `kill -INT`, because Ctrl-C's reaches rbspy directly.
+
+Send SIGTERM to `singed` alone, though, not to its whole process group as `kill -- -<pgid>` and `timeout` without `--foreground` do: sudo passes a SIGTERM straight on to rbspy, which then exits without writing the flamegraph. And rbspy kills only the command itself, so processes the command started may be left running.
+
 
 ## Limitations
 

lib/singed/cli.rb

@@ -22,6 +22,7 @@ class CLI
     def initialize(argv)
       @argv = argv
       @opts = OptionParser.new #: OptionParser
+      @interrupted = false #: bool
 
       parse_argv!
     end
@@ -87,12 +88,13 @@ def run
       )
       @filename = Singed::Flamegraph.generate_filename(label: "cli")
 
+      # nil values are for flags. rbspy uses its default rate unless one was given.
       options = {
         format: "speedscope",
         file: filename.to_s,
-        rate: @rate,
         silent: nil,
       }
+      options[:rate] = @rate if @rate
 
       rbspy_args = [
         "record",
@@ -163,7 +165,7 @@ def show_help?
       @show_help
     end
 
-    # Never nil or false: exception: true makes Kernel#system raise instead.
+    # Never nil or false: a command that fails raises instead.
     #: (Array[String | Integer | Pathname | nil], reason: String, ?env: Hash[String, String]) -> bool
     def sudo(system_args, reason:, env: {})
       loop do
@@ -182,8 +184,45 @@ def sudo(system_args, reason:, env: {})
       puts "$ #{Shellwords.join(sudo_args)}"
 
       # Sorbet can't check a splat of an array of unknown length: https://srb.help/7019
-      #: self as untyped
-      system(env, *sudo_args, exception: true)
+      process = Process #: as untyped
+      pid = process.spawn(env, *sudo_args) #: as Integer
+      status = wait_passing_on_signals(pid)
+      raise "#{Shellwords.join(sudo_args)} failed (#{status})" unless status.success?
+
+      true
+    end
+
+    # Kernel#system would leave the command running when singed is killed. Instead, the first SIGTERM is
+    # passed on as SIGINT, which sudo relays and is the only signal rbspy stops cleanly on, killing the
+    # command and writing the flamegraph. Nothing more is passed on, because rbspy exits without writing
+    # anything when interrupted twice, and Ctrl-C at a terminal already reaches it directly. Later SIGTERMs
+    # are ignored until singed exits, so they can't stop it opening the flamegraph either.
+    # https://github.com/rbspy/rbspy/blob/v0.53.0/src/main.rs#L164-L211
+    #: (Integer) -> Process::Status
+    def wait_passing_on_signals(pid)
+      previous_int = trap("INT", "IGNORE")
+      previous_term = trap("TERM") do
+        @interrupted ||= interrupt(pid)
+      end
+
+      # Process.wait2 only returns nil when told not to block.
+      waited = Process.wait2(pid) #: as !nil
+      waited.last
+    ensure
+      # trap returns nil for a handler installed outside Ruby, and restoring nil would ignore the signal.
+      trap("INT", previous_int || "DEFAULT")
+      trap("TERM", previous_term || "DEFAULT") unless @interrupted
+    end
+
+    # sudo before 1.9.13 doesn't relay a signal sent from its own process group, which singed shares to keep
+    # sudo in the terminal's foreground, so a kill in a group of its own sends it. kill fails while sudo
+    # briefly runs entirely as root, as it does starting up, and then a later SIGTERM tries again.
+    # https://github.com/sudo-project/sudo/commit/36742deec3041413af9b293706a64530321b96b5
+    #: (Integer) -> bool
+    def interrupt(pid)
+      kill = Process.spawn("kill", "-INT", pid.to_s, pgroup: true, err: File::NULL)
+      waited = Process.wait2(kill) #: as !nil
+      !!waited.last.success?
     end
 
     #: () -> String?

singed.gemspec

@@ -3,10 +3,10 @@
 Gem::Specification.new do |spec|
   spec.name = "singed"
 
-  spec.version = "0.3.0"
+  spec.version = "0.4.0"
   spec.license = "MIT"
-  spec.authors = ["Josh Nichols"]
-  spec.email = ["josh.nichols@gusto.com"]
+  spec.authors = ['Gusto Engineers']
+  spec.email = ['dev@gusto.com']
   spec.summary = "Quick and easy way to get flamegraphs from a specific part of your code base"
   spec.required_ruby_version = ">= 3.3"
   spec.homepage = "https://github.com/rubyatscale/singed"

spec/singed/cli_spec.rb

@@ -0,0 +1,192 @@
+# typed: false
+# frozen_string_literal: true
+
+require "rbconfig"
+require "singed/cli"
+
+# Runs exe/singed with stand-ins for sudo, for rbspy, and for the commands that open flamegraphs.
+RSpec.describe Singed::CLI do
+  let(:dir) { Pathname(Dir.mktmpdir("singed-cli-spec")) }
+  let(:bin) { dir.join("bin").tap(&:mkpath) }
+  let(:hold_open) { dir.join("hold_open") } # while it exists, opening the flamegraph doesn't finish
+  let(:interrupts) { dir.join("interrupts.log") } # a line for each SIGINT rbspy gets
+  let(:opened) { dir.join("opened.log") }
+  let(:options) { [] } # singed's, before the command
+  let(:output) { dir.join("output.log") }
+  let(:rbspy_args) { dir.join("rbspy_args.json") }
+  let(:rbspy_exit_status) { nil } # for rbspy to fail with, straight away
+  let(:started) { dir.join("started") } # the profiled command's pid, once it runs
+
+  # Writes the stand-ins first, so no path through the hooks runs singed with the real sudo. It runs as
+  # `bundle exec singed` would, but without this process's Bundler environment, whose BUNDLER_ORIG_PATH
+  # would have singed's Bundler.with_unbundled_env take the stand-ins back off PATH. And it runs in its
+  # own process group, so the after hook can clean up whatever singed leaves running. TMPDIR keeps the
+  # page that the bundled speedscope writes to show the flamegraph in dir too.
+  let!(:singed) do
+    write_stand_ins
+    Bundler.with_unbundled_env do
+      Process.spawn(
+        {
+          "PATH" => "#{bin}:#{ENV.fetch('PATH')}",
+          "BUNDLE_GEMFILE" => File.expand_path("../../Gemfile", __dir__),
+          "RUBYOPT" => "-rbundler/setup",
+          "TMPDIR" => dir.to_s,
+        },
+        RbConfig.ruby, File.expand_path("../../exe/singed", __dir__), "--output-directory", dir.to_s, *options,
+        "--", RbConfig.ruby, "-e", "File.write(ARGV[0], Process.pid.to_s); sleep", started.to_s,
+        chdir: dir.to_s, out: output.to_s, err: output.to_s, pgroup: true
+      )
+    end
+  end
+
+  after do
+    Process.kill("KILL", -singed)
+  rescue Errno::ESRCH, Errno::EPERM
+    # Everything has exited. macOS says EPERM when only an unreaped singed is left.
+  ensure
+    FileUtils.rm_rf(dir)
+  end
+
+  def write_stand_ins
+    # Like sudo before 1.9.13, relays SIGINT and SIGTERM that another process sends, but not from a process
+    # in its own process group that's still running. It's Perl because Ruby can't tell who sent a signal.
+    # https://github.com/sudo-project/sudo/blob/SUDO_1_9_12p2/src/exec_nopty.c#L150-L170
+    write_executable "sudo", <<~'PERL'
+      #!/usr/bin/env perl
+      use strict;
+      use warnings;
+      use POSIX ();
+
+      shift @ARGV while @ARGV && $ARGV[0] =~ /^-/;
+      my $command;
+      for my $signal (POSIX::SIGINT, POSIX::SIGTERM) {
+        my $relay = sub {
+          my $sender = $_[1]{pid} or return;
+          kill $signal, $command if $command && getpgrp($sender) != getpgrp(0);
+        };
+        POSIX::sigaction($signal, POSIX::SigAction->new($relay, POSIX::SigSet->new, POSIX::SA_SIGINFO));
+      }
+      $command = fork // die "fork: $!";
+      exec { $ARGV[0] } @ARGV or die "exec: $!" unless $command;
+      1 until waitpid($command, 0) == $command;
+      exit($? & 127 ? 128 + ($? & 127) : $? >> 8);
+    PERL
+    # Like rbspy, stops at the first SIGINT, then takes a moment to write the flamegraph, and exits
+    # without writing anything at a second SIGINT.
+    write_executable "rbspy", <<~RUBY
+      #!#{RbConfig.ruby}
+      #{"exit #{rbspy_exit_status}" if rbspy_exit_status}
+      require "json"
+      File.write(#{rbspy_args.to_s.inspect}, JSON.generate(ARGV))
+      file = ARGV[ARGV.index("--file") + 1]
+      trap("INT") do
+        File.write(#{interrupts.to_s.inspect}, "INT\\n", mode: "a")
+        exit!(1) if $interrupted
+        $interrupted = true
+        Process.kill("KILL", $command)
+      end
+      $command = Process.spawn(*ARGV.drop(ARGV.index("--") + 1))
+      Process.wait($command)
+      sleep 0.5 if $interrupted
+      File.write(file, JSON.generate(shared: { frames: [{ name: "<main>", file: "script.rb" }] }, profiles: []))
+    RUBY
+    ["npx", "open", "xdg-open"].each do |opener|
+      write_executable opener, <<~SH
+        #!/bin/sh
+        echo "$0 $*" >> #{opened}
+        while [ -e #{hold_open} ]; do sleep 0.01; done
+      SH
+    end
+  end
+
+  def write_executable(name, script)
+    bin.join(name).write(script)
+    bin.join(name).chmod(0o755)
+  end
+
+  def eventually(timeout: 60)
+    deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout
+    until (result = yield)
+      raise "Timed out. singed's output:\n#{output.read}" if Process.clock_gettime(Process::CLOCK_MONOTONIC) > deadline
+
+      sleep 0.02
+    end
+    result
+  end
+
+  def command_pid
+    eventually { started.exist? && started.read.to_i.nonzero? }
+  end
+
+  def exit_status
+    eventually { Process.wait2(singed, Process::WNOHANG)&.last }
+  end
+
+  it "stops rbspy and the profiled command when terminated, then opens the flamegraph" do
+    command = command_pid
+    Process.kill("TERM", singed)
+
+    expect(exit_status).to be_success, output.read
+    expect(interrupts.read).to eq("INT\n")
+    expect(JSON.parse(rbspy_args.read)).to start_with("record", "--format", "speedscope", "--file", a_string_ending_with(".json"), "--silent", "--")
+    expect { Process.kill(0, command) }.to raise_error(Errno::ESRCH)
+    expect(dir.glob("speedscope-cli-*.json")).not_to be_empty
+    expect(opened.read).not_to be_empty
+  end
+
+  it "passes on only the first SIGTERM, so rbspy can finish writing the flamegraph" do
+    command_pid
+    Process.kill("TERM", singed)
+    eventually { interrupts.exist? }
+    Process.kill("TERM", singed)
+
+    expect(exit_status).to be_success, output.read
+    expect(interrupts.read).to eq("INT\n")
+    expect(dir.glob("speedscope-cli-*.json")).not_to be_empty
+  end
+
+  it "ignores later SIGTERMs until it has opened the flamegraph" do
+    command_pid
+    FileUtils.touch(hold_open)
+    Process.kill("TERM", singed)
+    eventually { opened.exist? }
+    Process.kill("TERM", singed)
+    hold_open.delete
+
+    expect(exit_status).to be_success, output.read
+  end
+
+  it "leaves Ctrl-C's SIGINT to reach rbspy from the terminal" do
+    command_pid
+    Process.kill("INT", singed)
+    sleep 0.5
+
+    expect(Process.wait2(singed, Process::WNOHANG)).to be_nil
+    expect(interrupts).not_to exist
+
+    Process.kill("TERM", singed)
+
+    expect(exit_status).to be_success, output.read
+    expect(interrupts.read).to eq("INT\n")
+  end
+
+  context "with a rate" do
+    let(:options) { ["--rate", "50"] }
+
+    it "passes it on to rbspy" do
+      command_pid
+
+      expect(JSON.parse(rbspy_args.read)).to start_with("record", "--format", "speedscope", "--file", a_string_ending_with(".json"), "--silent", "--rate", "50", "--")
+    end
+  end
+
+  context "when rbspy fails" do
+    let(:rbspy_exit_status) { 3 }
+
+    it "fails too, without opening anything" do
+      expect(exit_status).not_to be_success
+      expect(output.read).to match(/rbspy record .* failed \(pid \d+ exit 3\)/)
+      expect(opened).not_to exist
+    end
+  end
+end