perftest

#!/usr/bin/ruby

DEFAULT_SYSTEM_CORE_COUNT = 2
USE_MLOCK_ALL = false

exit if ARGV.empty?
raise "/usr/bin/setarch not found" unless File.exist? "/usr/bin/setarch"
raise "/usr/bin/schedtool not found" unless File.exist? "/usr/bin/schedtool"

mlockall_source = <<-EOS
#include <stdio.h>
#include <sys/mman.h>

static void __attribute__((constructor)) ctor(void) {
    if( mlockall( MCL_CURRENT | MCL_FUTURE ) ) {
        perror( "mlockall failed" );
    }
}
EOS

if USE_MLOCK_ALL
    unless File.exist? "/tmp/libmlockall.so"
        `gcc -v &> /dev/null`
        raise "gcc not found" unless $?.exitstatus == 0

        File.write "/tmp/libmlockall.c", mlockall_source
        system "gcc -O2 -s -fPIC -shared /tmp/libmlockall.c -o /tmp/libmlockall.so"
        raise unless $?.exitstatus == 0
    end
end

original_governors = Dir["/sys/devices/system/cpu/cpu*/cpufreq/scaling_governor"].sort.map do |path|
    begin
        [path, File.read( path )]
        rescue Errno::EBUSY
            nil
    end
end.select { |entry| entry != nil }.to_h

cores = {}
File.read( "/proc/cpuinfo" ).strip.split("\n\n").each do |chunk|
    raise unless chunk =~ /^processor\s*:\s*(\d+)/
    nth_logical = $1.to_i
    raise unless chunk =~ /^core id\s*:\s*(\d+)/
    nth_physical = $1.to_i
    cores[nth_physical] ||= []
    cores[nth_physical] << nth_logical
end

raise if cores.empty?
raise "systems with a single core are not supported" if cores.length == 1

benchmark_cores = cores.dup
system_cores = {}

cores.take([DEFAULT_SYSTEM_CORE_COUNT, cores.length - 1].min).each do |nth, cpu|
    system_cores[nth] = cpu
    benchmark_cores.delete nth
end

benchmark_cpus_online = []
benchmark_cpus_offline = []

benchmark_cores.each do |_, cpus|
    benchmark_cpus_online << cpus[0]
    benchmark_cpus_offline += cpus[1..]
end

benchmark_cpus_offline_paths = benchmark_cpus_offline.map do |cpu|
    "/sys/devices/system/cpu/cpu#{cpu}/online"
end

system_cpu_mask = system_cores.values.flatten.sort.join(",")
benchmark_cpu_mask = benchmark_cpus_online.sort.join(",")

raise if benchmark_cpu_mask.empty?

pids = []
Dir["/proc/*"].each do |path|
    pid = File.basename( path )
    next unless pid =~ /\A\d+\Z/
    pid = pid.to_i
    pids << pid
end

original_affinities = {}

pids = pids.sort
pids.each do |pid|
    cgroup_path = "/proc/#{pid}/cgroup"
    if File.exist? cgroup_path
        begin
            data = File.read cgroup_path
            next if data.strip == "0::/"

            rescue Errno::ENOENT
                next
        end
    end

    output = `schedtool #{pid}`
    if $?.exitstatus == 0
        raise "unhandled output from schedtool: '#{output}'" unless output =~ /AFFINITY (0x[0-9a-f]+)/
        original_affinities[pid] = $1
    end
end

Signal.trap("INT") {}

r, w = IO.pipe
child_pid = fork do
    ld_preload = ENV["LD_PRELOAD"]
    if USE_MLOCK_ALL
        ld_preload += ":" unless ld_preload.empty?
        ld_preload += "/tmp/libmlockall.so"
    end
    ENV["LD_PRELOAD"] = ld_preload unless ld_preload == nil || ld_preload.empty?

    Process.kill "STOP", Process.pid
    w.write("0")
    exec *(["setarch", "-R"] + ARGV)
    exit
end

system "echo performance | sudo tee #{original_governors.keys.join(" ")} > /dev/null"
if $?.exitstatus != 0
    STDERR.puts "WARNING: failed to set the CPU governor to performance"
end

system "echo 0 | sudo tee #{benchmark_cpus_offline_paths.join(" ")} > /dev/null"
if $?.exitstatus != 0
    STDERR.puts "WARNING: failed to bring CPUs offline"
end

unless system_cpu_mask.empty?
    original_affinities.keys.each do |pid|
        system "sudo schedtool -a #{system_cpu_mask} #{pid} > /dev/null"
    end
end

system "sudo schedtool -a #{benchmark_cpu_mask} -n -20 #{child_pid}"
if $?.exitstatus != 0
    STDERR.puts "WARNING: schedtool failed"
end

if File.exist? "/sys/devices/system/cpu/cpufreq/boost"
    system "echo 0 | sudo tee /sys/devices/system/cpu/cpufreq/boost > /dev/null"
    if $?.exitstatus != 0
        STDERR.puts "WARNING: failed to disable turbo boost"
    end
end

loop do
    sleep 0.01
    Process.kill "CONT", child_pid
    begin
        r.read_nonblock 1
        rescue IO::EAGAINWaitReadable
            next
        else
            break
    end
end

Process.wait child_pid
status = $?.exitstatus

system "echo 1 | sudo tee #{benchmark_cpus_offline_paths.join(" ")} > /dev/null"
if $?.exitstatus != 0
    STDERR.puts "WARNING: failed to bring CPUs online"
end

if File.exist? "/sys/devices/system/cpu/cpufreq/boost"
    system "echo 1 | sudo tee /sys/devices/system/cpu/cpufreq/boost > /dev/null"
    if $?.exitstatus != 0
        STDERR.puts "WARNING: failed to reenable turbo boost"
    end
end

original_affinities.each do |pid, affinity|
    system "sudo schedtool -a #{affinity} #{pid} > /dev/null"
end

original_governors.each do |path, governor|
    system "echo '#{governor}' | sudo tee #{path} > /dev/null"
end

if status == nil
    exit 255
else
    exit status
end