#!/usr/bin/env bash
# Agent-session admission: at most AGENT_SESSION_SLOTS sessions run at once, a launch above
# the cap waits and is admitted when a holder exits, no session is ever signalled, and the
# session pool never gates the builds that run inside a session.
set -uo pipefail

MODULE="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
TESTROOT=$(mktemp -d "${TMPDIR:-/tmp}/agent-session-cap.XXXXXX")
PIDS=()
cleanup() {
  local pid
  for pid in "${PIDS[@]:-}"; do
    [[ -n "$pid" ]] || continue
    kill "$pid" 2>/dev/null
  done
  rm -rf "$TESTROOT"
}
trap cleanup EXIT

BIN="$TESTROOT/bin"
mkdir -p "$BIN" "$TESTROOT/lib" "$TESTROOT/session-slots" "$TESTROOT/build-slots"
cp "$MODULE/bin/_agent-session-admission" "$BIN/_agent-session-admission"
cp "$MODULE/lib/buildslot.sh" "$TESTROOT/lib/buildslot.sh"
# The real scope script would put every session in a systemd scope; the cap is what is
# under test, so the chain past admission is reduced to an exec.
printf '#!/usr/bin/env bash\nexec "$@"\n' >"$BIN/_agent-build-scope"
chmod +x "$BIN/_agent-build-scope" "$BIN/_agent-session-admission"

pass=0
fail=0
ok() { printf 'ok   %s\n' "$1"; pass=$((pass + 1)); }
bad() { printf 'FAIL %s\n' "$1"; fail=$((fail + 1)); }
check() { if [[ "$1" == "$2" ]]; then ok "$3"; else bad "$3 (want '$2', got '$1')"; fi; }

# A session that reports it started and then runs until its own release file appears.
session() {
  local name=$1
  shift
  env "$@" AGENT_SESSION_SLOTS=2 AGENT_SESSION_SLOT_DIR="$TESTROOT/session-slots" \
    "$BIN/_agent-session-admission" /bin/bash -c '
      printf started >"$1.started"
      while [[ ! -e "$1.release" ]]; do sleep 0.05; done
      printf done >"$1.exited"
    ' _ "$TESTROOT/$name" &
  PIDS+=("$!")
}

await() {
  local path=$1 deadline=$((SECONDS + ${2:-20}))
  while (( SECONDS < deadline )); do
    [[ -e "$path" ]] && return 0
    sleep 0.05
  done
  return 1
}

session first
session second
if await "$TESTROOT/first.started" && await "$TESTROOT/second.started"; then
  ok "sessions up to the cap start immediately"
else
  bad "sessions up to the cap start immediately"
fi

# BUILD_SLOT_TIMEOUT is set on the ambient environment on purpose: a build gate that asked
# for a bounded wait must not turn a capped session launch into a failed launch.
session third BUILD_SLOT_TIMEOUT=1
sleep 3
if [[ ! -e "$TESTROOT/third.started" ]] && kill -0 "${PIDS[2]}" 2>/dev/null; then
  ok "a session above the cap waits instead of starting or failing"
else
  bad "a session above the cap waits instead of starting or failing"
fi

for pid in "${PIDS[0]}" "${PIDS[1]}"; do
  if kill -0 "$pid" 2>/dev/null; then
    ok "running session $pid is untouched while another waits"
  else
    bad "running session $pid is untouched while another waits"
  fi
done

: >"$TESTROOT/first.release"
if await "$TESTROOT/third.started"; then
  ok "the waiting session is admitted once a holder exits"
else
  bad "the waiting session is admitted once a holder exits"
fi
check "$([[ -e "$TESTROOT/first.exited" ]] && printf yes || printf no)" yes \
  "the released session exited on its own"
if kill -0 "${PIDS[1]}" 2>/dev/null; then
  ok "the untouched session is still running after the queue drained"
else
  bad "the untouched session is still running after the queue drained"
fi

: >"$TESTROOT/second.release"
: >"$TESTROOT/third.release"
wait "${PIDS[0]}" "${PIDS[1]}" "${PIDS[2]}" 2>/dev/null

# What the session hands to the code running inside it.
inside=$(AGENT_SESSION_SLOTS=2 AGENT_SESSION_SLOT_DIR="$TESTROOT/session-slots" \
  BUILD_SLOT_DIR="$TESTROOT/build-slots" BUILD_SLOTS=7 BUILD_SLOT_TIMEOUT=1 \
  "$BIN/_agent-session-admission" /bin/bash -c \
  'printf "%s|%s|%s|%s" "${BUILD_SLOT_HELD:-unset}" "${BUILD_SLOT_DIR:-unset}" "${BUILD_SLOTS:-unset}" "${BUILD_SLOT_ISOLATED:-unset}"')
check "$inside" "unset|unset|unset|unset" "a session inherits neither the slot nor the pool it was admitted from"

# One session holds a session slot; a build launched inside it must still face the build
# gate, and the build pool must be untouched by the session pool.
session fourth
await "$TESTROOT/fourth.started" || bad "fourth session started"
build_rc=$(BUILD_SLOT_DIR="$TESTROOT/build-slots" BUILD_SLOTS=1 BUILD_SLOT_FIXED=1 \
  "$TESTROOT/lib/buildslot.sh" /bin/true >/dev/null 2>&1; printf '%s' $?)
check "$build_rc" 0 "a build is admitted while a session holds a session slot"
held=$(find "$TESTROOT/build-slots" -name 'slot-*.lock' | wc -l)
check "$held" 1 "the session took no slot from the build pool"
: >"$TESTROOT/fourth.release"
wait "${PIDS[3]}" 2>/dev/null

# Nested launch: a runtime started inside a session's process tree must never queue behind
# its own ancestor.
session fifth
await "$TESTROOT/fifth.started" || bad "fifth session started"
nested=$(AGENT_SESSION_SLOTS=2 AGENT_SESSION_SLOT_DIR="$TESTROOT/session-slots" \
  "$BIN/_agent-session-admission" /bin/bash -c '
    timeout 5 env AGENT_SESSION_SLOTS=2 AGENT_SESSION_SLOT_DIR="$1" \
      "$2/_agent-session-admission" /bin/echo nested
  ' _ "$TESTROOT/session-slots" "$BIN" 2>/dev/null)
check "$nested" nested "a nested session inherits its ancestor's slot instead of deadlocking"
: >"$TESTROOT/fifth.release"
wait "${PIDS[4]}" 2>/dev/null

printf 'agent-session-cap: pass=%d fail=%d\n' "$pass" "$fail"
(( fail == 0 ))