#!/bin/bash
# 실습 #4 ⑤ — VM 친화 부하 발생 (stress-ng 의존성 없음, 표준 도구만)
# VirtualBox / VMware / 클라우드 VM 에서 AVX SIGILL 없이 100% 동작 보장.
#
# 부하 종류별 도구:
#   CPU  → yes > /dev/null & (단순 루프) · openssl speed (FPU·SSE 만)
#   MEM  → python3 (제어 가능한 양 + 즉시 해제)
#   IO   → dd + fsync (VM 디스크 캐시 무시)
#   NET  → iperf3 loopback (VM 안에서 자가 측정)
#   LOAD → 다중 sha256sum (load average 인위적 상승)
#
# 사용:
#   bash setup_05_vm_safe_load.sh menu        # 인터랙티브 메뉴
#   bash setup_05_vm_safe_load.sh cpu 30      # CPU 부하 30초
#   bash setup_05_vm_safe_load.sh mem 1024 10 # 메모리 1024MB × 10초
#   bash setup_05_vm_safe_load.sh io          # 디스크 I/O 폭풍
#   bash setup_05_vm_safe_load.sh load 8 30   # load avg 인위적 8 × 30초
#   bash setup_05_vm_safe_load.sh stop        # 모든 백그라운드 부하 종료
set -e

ACTION="${1:-menu}"
WORK=/tmp/lab04_vmsafe
mkdir -p "$WORK"
PIDFILE="$WORK/load.pids"

# ─── 공용 ──────────────────────────────────────────────
stop_all() {
    if [ -f "$PIDFILE" ]; then
        while read -r p; do
            kill "$p" 2>/dev/null || true
        done < "$PIDFILE"
        rm -f "$PIDFILE"
    fi
    # yes / sha256sum 잔류 정리
    pkill -f "yes > /dev/null" 2>/dev/null || true
    pkill -f "sha256sum /dev/zero" 2>/dev/null || true
    echo "✓ 모든 부하 종료"
}

# ─── CPU 부하 (yes 루프, N 코어) ──────────────────────
cpu_load() {
    DUR="${1:-30}"
    CORES=$(nproc)
    echo "🔥 CPU 부하 — yes loop × $CORES 코어 × ${DUR}초"
    >"$PIDFILE"
    for i in $(seq 1 $CORES); do
        ( yes > /dev/null ) &
        echo $! >> "$PIDFILE"
    done
    sleep "$DUR"
    stop_all
}

# ─── 메모리 부하 (Python — 정확한 양, 즉시 해제) ───────
mem_load() {
    SIZE_MB="${1:-512}"
    DUR="${2:-10}"
    echo "🔥 메모리 부하 — ${SIZE_MB}MB × ${DUR}초"
    python3 -c "
import time, sys
n = $SIZE_MB
buf = bytearray(n * 1024 * 1024)
print(f'  Allocated {n} MB')
sys.stdout.flush()
time.sleep($DUR)
print('  Released')
" &
    echo $! > "$PIDFILE"
    wait $(cat "$PIDFILE")
    rm -f "$PIDFILE"
}

# ─── 디스크 I/O (dd + fsync — VM 디스크 캐시 우회) ───
io_load() {
    SIZE="${1:-500}"        # MB
    FILE="$WORK/io_test.bin"
    echo "🔥 디스크 쓰기 — ${SIZE}MB (fsync 강제)"
    dd if=/dev/zero of="$FILE" bs=1M count="$SIZE" conv=fsync status=progress
    echo
    echo "🔥 디스크 읽기 — ${SIZE}MB"
    dd if="$FILE" of=/dev/null bs=1M
    rm -f "$FILE"
    echo "✓ 검증: iostat -xz 1 의 %util 컬럼이 솟구쳐야 함"
}

# ─── 네트워크 부하 (iperf3 loopback) ───────────────────
net_load() {
    if ! command -v iperf3 >/dev/null; then
        echo "→ iperf3 설치 필요: sudo apt install -y iperf3"
        return 1
    fi
    DUR="${1:-10}"
    echo "🔥 네트워크 (loopback) — ${DUR}초"
    iperf3 -s -D --logfile "$WORK/iperf-srv.log"
    sleep 1
    iperf3 -c 127.0.0.1 -t "$DUR"
    pkill iperf3 2>/dev/null || true
}

# ─── load average 인위적 상승 (sha256sum × N) ─────────
load_avg() {
    N="${1:-8}"
    DUR="${2:-30}"
    echo "🔥 load average 부하 — sha256sum × $N × ${DUR}초"
    echo "→ 다른 터미널에서 uptime 으로 load avg 1m 컬럼 관찰"
    >"$PIDFILE"
    for i in $(seq 1 $N); do
        ( sha256sum /dev/zero > /dev/null ) &
        echo $! >> "$PIDFILE"
    done
    sleep "$DUR"
    stop_all
}

# ─── 종합 4축 부하 (CPU+MEM+IO 동시) ─────────────────
all_load() {
    DUR="${1:-30}"
    echo "🔥 4축 종합 — CPU + MEM + IO × ${DUR}초"
    # CPU
    >"$PIDFILE"
    for i in $(seq 1 $(nproc)); do ( yes > /dev/null ) & echo $! >> "$PIDFILE"; done
    # IO (백그라운드)
    ( dd if=/dev/zero of="$WORK/io_test.bin" bs=1M count=500 conv=fsync 2>/dev/null \
      && rm -f "$WORK/io_test.bin" ) &
    echo $! >> "$PIDFILE"
    # MEM (백그라운드)
    ( python3 -c "
import time
buf = bytearray(512 * 1024 * 1024)
time.sleep($DUR)
" ) &
    echo $! >> "$PIDFILE"
    sleep "$DUR"
    stop_all
}

# ─── 인터랙티브 메뉴 ────────────────────────────────
menu() {
    cat <<MSG
═══ VM 친화 부하 메뉴 ═══

  1) CPU       (yes 루프 × N 코어)              ← 가장 안전
  2) Memory    (Python — 정확한 양)
  3) Disk I/O  (dd + fsync)
  4) Network   (iperf3 loopback)
  5) Load avg  (sha256sum × N)
  6) 4축 종합  (CPU + MEM + IO)
  s) 전체 종료
  q) 나가기

옆 터미널에서 동시 관찰: htop / vmstat 1 / iostat -xz 1 / uptime
MSG
    read -rp "선택: " ch
    case "$ch" in
        1) read -rp "지속 초 (기본 30): " s; cpu_load "${s:-30}" ;;
        2) read -rp "크기 MB (기본 512): " m; read -rp "초 (기본 10): " s; mem_load "${m:-512}" "${s:-10}" ;;
        3) io_load 500 ;;
        4) net_load 10 ;;
        5) read -rp "프로세스 수 (기본 8): " n; read -rp "초 (기본 30): " s; load_avg "${n:-8}" "${s:-30}" ;;
        6) all_load 30 ;;
        s) stop_all ;;
        *) echo "취소" ;;
    esac
}

# ─── 진입점 ───────────────────────────────────────
case "$ACTION" in
    cpu)  cpu_load  "${2:-30}" ;;
    mem)  mem_load  "${2:-512}" "${3:-10}" ;;
    io)   io_load   "${2:-500}" ;;
    net)  net_load  "${2:-10}" ;;
    load) load_avg  "${2:-8}"   "${3:-30}" ;;
    all)  all_load  "${2:-30}" ;;
    stop) stop_all ;;
    menu) menu ;;
    *)
        echo "사용법:"
        echo "  bash $0 menu"
        echo "  bash $0 {cpu|mem|io|net|load|all|stop} [인자...]"
        exit 1 ;;
esac
