xvfb-run

#!/bin/sh
#
# Robust xvfb-run wrapper for headless / HPC environments
#
# Key features:
#  - Clear error logging
#  - No use of `tempfile`, `$$`, or explicit XXXXX templates
#  - Falls back to running command directly if Xvfb is unavailable
#  - Forces headless plotting backends when falling back
#

set -u

PROGNAME="$(basename "$0")"

SERVERNUM=""
AUTO_SERVERNUM=""
ERRORFILE=""
XVFBARGS="-screen 0 1280x1024x24"
XVFB_RUN_TMPDIR=""
AUTHFILE=""
XVFBPID=""

XAUTHORITY_ORIG="${XAUTHORITY-}"
DISPLAY_ORIG="${DISPLAY-}"

usage() {
    cat >&2 <<EOF
Usage: ${PROGNAME} [OPTIONS] command [args...]

Options:
  -a            Automatically choose a free X server number (starting at 99)
  -n NUM        Use X server number NUM
  -e FILE       Append logs to FILE (default: stderr only)
  -s ARGS       Extra arguments for Xvfb
  -h            Show this help

Behavior:
  If Xvfb is not available on PATH, the command is run directly with
  headless backends forced (MPLBACKEND=Agg, QT_QPA_PLATFORM=offscreen).
EOF
}

log() {
    msg="[$PROGNAME] $*"
    echo "$msg" >&2
    if [ -n "$ERRORFILE" ]; then
        echo "$msg" >>"$ERRORFILE" 2>/dev/null || true
    fi
}

die() {
    log "ERROR: $*"
    exit 1
}

cleanup() {
    if [ -n "$XVFBPID" ] && kill -0 "$XVFBPID" 2>/dev/null; then
        kill "$XVFBPID" 2>/dev/null || true
        wait "$XVFBPID" 2>/dev/null || true
    fi
    if [ -n "$XVFB_RUN_TMPDIR" ] && [ -d "$XVFB_RUN_TMPDIR" ]; then
        rm -rf "$XVFB_RUN_TMPDIR" 2>/dev/null || true
    fi
    if [ -n "$XAUTHORITY_ORIG" ]; then
        export XAUTHORITY="$XAUTHORITY_ORIG"
    else
        unset XAUTHORITY 2>/dev/null || true
    fi
    if [ -n "$DISPLAY_ORIG" ]; then
        export DISPLAY="$DISPLAY_ORIG"
    else
        unset DISPLAY 2>/dev/null || true
    fi
}
trap cleanup EXIT HUP INT TERM

# Parse options
while getopts "an:e:s:h" opt; do
    case "$opt" in
        a) AUTO_SERVERNUM=1 ;;
        n) SERVERNUM="$OPTARG" ;;
        e) ERRORFILE="$OPTARG" ;;
        s) XVFBARGS="$XVFBARGS $OPTARG" ;;
        h) usage; exit 0 ;;
        *) usage; exit 2 ;;
    esac
done
shift $((OPTIND - 1))

[ "$#" -ge 1 ] || { usage; exit 2; }

# --- Fallback path: no Xvfb available ---
if ! command -v Xvfb >/dev/null 2>&1; then
    log "WARNING: Xvfb not found; running command without virtual display."
    export MPLBACKEND="${MPLBACKEND-Agg}"
    export QT_QPA_PLATFORM="${QT_QPA_PLATFORM-offscreen}"
    exec "$@"
fi

# Required tools if we are using Xvfb
command -v mktemp >/dev/null 2>&1 || die "Required command not found: mktemp"
command -v xauth  >/dev/null 2>&1 || die "Required command not found: xauth"
command -v mcookie >/dev/null 2>&1 || die "Required command not found: mcookie"

# Choose server number
if [ -z "$SERVERNUM" ]; then
    SERVERNUM=99
    if [ -n "$AUTO_SERVERNUM" ]; then
        while [ -f "/tmp/.X${SERVERNUM}-lock" ]; do
            SERVERNUM=$((SERVERNUM + 1))
            [ "$SERVERNUM" -le 200 ] || die "Could not find free X server number"
        done
    fi
fi

# Temp dir and Xauthority
XVFB_RUN_TMPDIR="$(mktemp -d)" || die "mktemp -d failed"
AUTHFILE="$XVFB_RUN_TMPDIR/Xauthority"
: >"$AUTHFILE" || die "Cannot create $AUTHFILE"

COOKIE="$(mcookie)" || die "mcookie failed"
if ! XAUTHORITY="$AUTHFILE" xauth -q add ":$SERVERNUM" MIT-MAGIC-COOKIE-1 "$COOKIE" \
    2>>"${ERRORFILE:-/dev/stderr}"; then
    die "xauth failed"
fi

# Start Xvfb
log "Starting Xvfb on :$SERVERNUM"
(
    exec Xvfb ":$SERVERNUM" -auth "$AUTHFILE" $XVFBARGS
) >>"${ERRORFILE:-/dev/stderr}" 2>&1 &
XVFBPID=$!

# Wait for socket
i=0
while [ $i -lt 50 ]; do
    if [ -S "/tmp/.X11-unix/X$SERVERNUM" ]; then
        break
    fi
    if ! kill -0 "$XVFBPID" 2>/dev/null; then
        die "Xvfb died during startup"
    fi
    i=$((i + 1))
    sleep 0.05 2>/dev/null || sleep 1
done

[ -S "/tmp/.X11-unix/X$SERVERNUM" ] || die "Xvfb did not start"

# Run command
export DISPLAY=":$SERVERNUM"
export XAUTHORITY="$AUTHFILE"

log "Running: $*"
exec "$@"