cgexec-wrapper

#!/bin/bash

set -e
# This is a drop-in replacement for cgexec, with the following differences:
#
# - It will put child processes in a cgroup relative to THIS process's cgroup.

# Required arguments:
# -g <controllers>:<CgroupName>

# Function to display usage information
usage() {
    echo "Usage: $0 -g <value>" >&2
    exit 1
}

# Parse command line arguments
while getopts ":g:" opt; do
  case ${opt} in
    g )
      # Split the value on ":" and assign to variables
      IFS=':' read -r -a parts <<< "$OPTARG"
      controllers="${parts[0]}"
      cgroupName="${parts[1]}"
      ;;
    \? )
      usage
      ;;
  esac
done
shift $((OPTIND -1))

# Check if both variables are set, otherwise show an error message and exit
if [ -z "$controllers" ] || [ -z "$cgroupName" ]; then
    echo "Error: Both controllers and cgroupName must be specified." >&2
    usage
fi

# Get my own Cgroup from /proc/self/cgroup.
# There are three fields here, separated by ':':
# - the group ID. This is always 0 for cgroups v2.
# - the enabled controllers. This is always empty for cgroups v2.
# - the cgroup Path. This is the only part we care about.
cgroupNameRelativePath="$(cut -d':' -f3- < /proc/self/cgroup)"

# Buildfarm's Java worker evacuates itself to an "evacuation" subdirectory during
# initialization for cgroupsv2. The execution cgroups are created relative to the
# ORIGINAL container cgroup (the parent of "evacuation"), not the evacuation cgroup itself.
# If we're in an evacuation cgroup, use its parent as the base.
if [[ "$cgroupNameRelativePath" == */evacuation ]]; then
    cgroupBasePath=$(dirname "$cgroupNameRelativePath")
else
    cgroupBasePath="$cgroupNameRelativePath"
fi

# This diverges from `cgexec` where we place ourselves in a group relative to the base cgroup.
# Example:
# cgroupName = `executions/operations/29aec796-ad19-4e6b-a83b-b8b9f55d0abb`
# Worker process is in the "evacuation" subdirectory:
#   ${cgroupBasePath} (e.g. /kubepods/pod-xxx/container-yyy)
#   ├─ evacuation/
#   │  ├─ 123456 /tini -- java -jar /app/build_buildfarm/buildfarm-shard-worker_deploy.jar
#   │  └─ 123457 java -jar /app/build_buildfarm/buildfarm-shard-worker_deploy.jar
#   └─ executions/
#      └─ operations/
#         └─ 29aec796-ad19-4e6b-a83b-b8b9f55d0abb/
#            └─ 123458 ${@}

# Target cgroup.procs path
cgroupProcsPath="/sys/fs/cgroup${cgroupBasePath}/$cgroupName/cgroup.procs"

# Debug info on failure
if ! echo $$ > "$cgroupProcsPath" 2>/dev/null; then
    echo "cgexec-wrapper: Failed to write PID $$ to $cgroupProcsPath" >&2
    echo "cgexec-wrapper: Current user: $(id)" >&2
    echo "cgexec-wrapper: cgroup.procs permissions: $(ls -la "$cgroupProcsPath" 2>&1)" >&2
    echo "cgexec-wrapper: Parent dir: $(ls -la "$(dirname "$cgroupProcsPath")" 2>&1 | head -5)" >&2
    # Fail hard - do NOT allow the action to run without cgroup limits,
    # as it would run unconstrained and could OOM the entire container.
    echo "cgexec-wrapper: FATAL: Refusing to execute without cgroup limits" >&2
    exit 1
fi

# Exec the rest of the arguments as-is
exec "$@"