bench.sh

#! /usr/bin/env nix-shell
#! nix-shell -i bash -p bash coreutils procps cpuset which

# shellcheck shell=bash

# Based on suggestions from here: https://llvm.org/docs/Benchmarking.html.
# Also see: https://juliaci.github.io/BenchmarkTools.jl/dev/linuxtips/.
# And: https://man7.org/linux/man-pages/man7/cpuset.7.html.

set -euo pipefail

readonly BOLD='\033[1m'
readonly RED='\033[0;31m'
readonly NC='\033[0m'

echo() { command echo -e >&2 "[bench]" "${@}"; }
err() {
  echo "${RED}error${NC}: $*"
  exit "${ec-1}"
}
join_by() { local IFS="$1"; shift; command echo "$*"; }

if [[ $(id -u) != 0 ]]; then
  ec=2 err "please run as root"
fi

# $@: cores
check_cpu_governor() {
  local err=false
  for cpu in "$@"; do
    local gov="/sys/devices/system/cpu/cpu${cpu}/cpufreq/scaling_governor"
    local g
    g="$(cat "$gov")"
    if [[ "$g" != "performance" ]]; then
      echo "${RED}error${NC}: CPU $cpu has governor: '${RED}${g}${NC}' instead of 'performance'"
      err=true
    fi
  done

  if [[ $err == true ]]; then
    echo
    ec=3 err "not all CPUs set to 'performance' CPU governor"
  fi
}

check_aslr_disabled() {
  if [[ "$(sysctl kernel.randomize_va_space -n)" != "0" ]]; then
    # shellcheck disable=SC2016
    ec=4 err 'ASLR does not seem to be disabled; try `sysctl kernel.randomize_va_space=0`?'
  fi
}

check_turbo_disabled() {
  local turbo_disable_file="/sys/devices/system/cpu/intel_pstate/no_turbo"

  if ! [[ -f ${turbo_disable_file} ]]; then
    echo "warning: couldn't find no_turbo sysfs file.."
    return
  fi

  if [[ "$(cat /sys/devices/system/cpu/intel_pstate/no_turbo)" != "1" ]]; then
    ec=4 err 'turbo boost does not seem to be disabled; try `echo 1 >'" $turbo_disable_file"'`?'
  fi

}

# This is the list of logical cores mapped to corresponding sibling hyperthreads that we can disable:
declare -A logical_cores=()
declare total_available_cores
total_available_cores=$(nproc --all)
for t in /sys/devices/system/cpu/cpu*/topology/thread_siblings_list; do
  # shellcheck disable=SC2012
  core=$(ls "$t" | cut -d'/' -f 6 | cut -d 'u' -f2)
  first=$(<"$t" cut -d',' -f1)
  rest="$(<"$t" cut -d',' -f2-)"
  if [[ $first == "$core" ]]; then
     logical_cores["$core"]="$rest"
  fi
done

# $@: list of core numbers
declare CHECK_HT_DISABLED=true
check_cores() {
  local failed=false

  for core in "${@}"; do
    # check that the core is online:
    #
    # (CPU 0 is always online..)
    if [[ ${core} != 0 ]] && [[ "$(cat /sys/devices/system/cpu/cpu${core}/online)" != "1" ]]; then
      echo "${RED}error${NC}: CPU $core is not online!"
      failed=true
    fi

    if [[ $CHECK_HT_DISABLED == "true" ]]; then
      # check that all sibling hyperthreads are disabled:
      IFS="," read -ra sibling_hyperthreads <<<"${logical_cores[$core]}"
      for sibling in "${sibling_hyperthreads[@]}"; do
        if [[ $sibling == $core ]]; then continue; fi # `cut` is annoying..
        if [[ "$(cat "/sys/devices/system/cpu/cpu${sibling}/online")" != "0" ]]; then
          echo "${RED}error${NC}: hyperthread sibling (CPU $sibling) of CPU $core is online; expected it to be disabled!"
          failed=true
        fi
      done
    fi
  done

  if [[ $failed == "true" ]]; then
     ec=6 err "CPUs and hyperthreads not configured as expected"
  fi
}

cores=${#logical_cores[@]}
case "${1-""}" in
  -j*) cores="${1/-j}"; shift ;;
  *) ;;
esac
readonly cores

if [[ ${#} == 0 ]]; then
  ec=9 err "usage: [-jN] <command> [args...]"
fi

declare -a core_list
# We want to pick logical cores and exclude the first core, if possible:
if [[ ${cores} -lt "${#logical_cores[@]}" ]]; then
  # Fewer than the number of logical cores were requested; we can use
  # cores from `1 ... ${cores}`
  # shellcheck disable=SC2207
  core_list=($(seq 1 "${cores}"))
elif [[ ${cores} == "${#logical_cores[@]}" ]]; then
  # We're using all the logical cores; we have to use cores from:
  # `0 .. $((${cores} - 1))`
  # shellcheck disable=SC2207
  core_list=($(seq 0 $((cores-1))))
elif [[ ${cores} -le "${total_available_cores}" ]]; then
  # Asked for more threads than we have logical cores, need to use hyper-threads.
  CHECK_HT_DISABLED=false
  echo "warning: there are only ${#logical_cores[@]} logical cores present but $cores cores were requested; this will use ${BOLD}hyper-threads${NC}"
  # shellcheck disable=SC2207
  core_list=($(seq 0 $((cores-1))))
else
  # Asked for more threads that we have, outright.
  ec=6 err "there are only ${BOLD}${total_available_cores} threads${NC} available on this system but ${RED}${cores} cores${NC} were requested.."
fi

echo "info: using ${BOLD}$cores cores${NC} (out of ${total_available_cores} total available; ${#logical_cores[@]} logical available)"
echo "cores: ${core_list[*]@E}"

check_aslr_disabled
check_turbo_disabled
check_cpu_governor "${core_list[@]}"
check_cores "${core_list[@]}"

##################

readonly CSET_NAME="user-cset-$$"

command echo >&2
echo "setting up cpuset ${BOLD}$CSET_NAME${NC}:"
cset shield ${VERBOSE+-vvvvv} \
  --userset="$CSET_NAME" --sysset="system" \
  --cpu="$(join_by , "${core_list[@]}")" \
  --kthread=on >&2 || {
    ec=$? err 'cset setup error; try deleting cpusets with `sudo '"$(which cset) set -d system -d $CSET_NAME"'`'
}

cleanup() {
  command echo >&2
  echo "cleanup:"
  cset shield ${VERBOSE+-vvvvv} \
    --force \
    --userset="${CSET_NAME}" --sysset="system" \
    --reset >&2
}
trap cleanup EXIT

sigint_count=0
force_quit() {
  # just give some indication that we were interrupted:
  ec=10 err 'force quit'

  # we can't resume after being interrupted so this doesn't really work:
  case $((++sigint_count)) in
    1) command echo -e >&2 "${BOLD}Are you sure you want to quit?${NC}";;
    2) command echo -e >&2 "${BOLD}Last chance!${NC}";;
    *) ;;
  esac
}
trap force_quit SIGINT

command echo >&2
echo "running [under CSET: ${BOLD}$CSET_NAME${NC}]: ${BOLD}${*@Q}${NC}"
cset shield \
  --userset="$CSET_NAME" --sysset="system" \
  --exec -- \
  "${@}" || {
  exit_code=$?
}

exit ${exit_code-0}