#!/bin/bash
# 2025-05-24 add by dukai

set -euo pipefail

DEFAULT_PATH="/tmp"
DEFAULT_NUMJOBS=1
DEFAULT_RW="randread"
DEFAULT_BS="4k"
DEFAULT_IODEPTH=1
DEFAULT_SIZE="80G"
DEFAULT_RUNTIME=300
DEFAULT_RAMP_TIME=10

usage() {
    cat <<HELP
用法:
  $(basename "$0") [选项]

选项:
  -p, --path      <路径>      测试文件所在目录    （默认: ${DEFAULT_PATH}）
  -j, --numjobs   <并发数>    并发工作线程数      （默认: ${DEFAULT_NUMJOBS}）
  -r, --rw        <策略>      读写策略            （默认: ${DEFAULT_RW}）
  -b, --bs        <块大小>    I/O 块大小          （默认: ${DEFAULT_BS}）
  -d, --iodepth   <队列深度>  I/O 队列深度        （默认: ${DEFAULT_IODEPTH}）
  -s, --size      <大小>      测试文件大小        （默认: ${DEFAULT_SIZE}）
  -t, --runtime   <秒>        测试持续时间        （默认: ${DEFAULT_RUNTIME}s）
      --ramp      <秒>        预热时间            （默认: ${DEFAULT_RAMP_TIME}s）
  -o, --output    <目录>      结果输出目录        （默认: 当前目录）
  -h, --help                  显示此帮助

支持的 --rw 策略:
  read        顺序读
  write       顺序写
  randread    随机读  ← 默认
  randwrite   随机写
  readwrite   顺序混合读写（50/50）
  randrw      随机混合读写（50/50）
  trimwrite   Trim + 写（NVMe 专用）

示例:
  $(basename "$0") --path /data --numjobs 8 --rw randread --iodepth 32
HELP
    exit 0
}

PATH_ARG="$DEFAULT_PATH"
NUMJOBS="$DEFAULT_NUMJOBS"
RW="$DEFAULT_RW"
BS="$DEFAULT_BS"
IODEPTH="$DEFAULT_IODEPTH"
SIZE="$DEFAULT_SIZE"
RUNTIME="$DEFAULT_RUNTIME"
RAMP_TIME="$DEFAULT_RAMP_TIME"
OUTPUT_DIR="."

while [[ $# -gt 0 ]]; do
    case "$1" in
        -p|--path)     PATH_ARG="${2:?'--path 需要参数'}";    shift 2 ;;
        -j|--numjobs)  NUMJOBS="${2:?'--numjobs 需要参数'}";  shift 2 ;;
        -r|--rw)       RW="${2:?'--rw 需要参数'}";            shift 2 ;;
        -b|--bs)       BS="${2:?'--bs 需要参数'}";            shift 2 ;;
        -d|--iodepth)  IODEPTH="${2:?'--iodepth 需要参数'}";  shift 2 ;;
        -s|--size)     SIZE="${2:?'--size 需要参数'}";        shift 2 ;;
        -t|--runtime)  RUNTIME="${2:?'--runtime 需要参数'}";  shift 2 ;;
        --ramp)        RAMP_TIME="${2:?'--ramp 需要参数'}";   shift 2 ;;
        -o|--output)   OUTPUT_DIR="${2:?'--output 需要参数'}"; shift 2 ;;
        -h|--help)     usage ;;
        *) echo "[ERROR] 未知参数: $1，使用 --help 查看帮助" >&2; exit 1 ;;
    esac
done

VALID_RW="read write randread randwrite readwrite randrw trimwrite"
echo "$VALID_RW" | grep -qw "$RW" \
    || { echo "[ERROR] 不支持的 --rw 策略: '$RW'。支持: $VALID_RW" >&2; exit 1; }

[[ "$NUMJOBS"  =~ ^[1-9][0-9]*$ ]] || { echo "[ERROR] --numjobs 必须为正整数"  >&2; exit 1; }
[[ "$IODEPTH"  =~ ^[1-9][0-9]*$ ]] || { echo "[ERROR] --iodepth 必须为正整数"  >&2; exit 1; }
[[ "$RUNTIME"  =~ ^[1-9][0-9]*$ ]] || { echo "[ERROR] --runtime 必须为正整数"  >&2; exit 1; }

command -v fio &>/dev/null || { echo "[ERROR] 未找到 fio，请先安装: apt install fio / yum install fio" >&2; exit 1; }
[[ -d "$PATH_ARG" ]] || { echo "[ERROR] 测试路径不存在: $PATH_ARG" >&2; exit 1; }
[[ -w "$PATH_ARG" ]] || { echo "[ERROR] 测试路径无写权限: $PATH_ARG" >&2; exit 1; }
mkdir -p "$OUTPUT_DIR" || { echo "[ERROR] 无法创建输出目录: $OUTPUT_DIR" >&2; exit 1; }

TIMESTAMP=$(date '+%Y%m%d_%H%M%S')
JOB_NAME="${RW}_${BS}_j${NUMJOBS}_qd${IODEPTH}"
TEST_FILE="${PATH_ARG}/fio_test_${TIMESTAMP}.tmp"
OUTPUT_PREFIX="${OUTPUT_DIR}/${JOB_NAME}_${TIMESTAMP}"
SUMMARY_FILE="${OUTPUT_PREFIX}_summary.txt"

cat <<CONF

================================================
  fio 磁盘 I/O 性能测试
================================================
  测试路径  : $PATH_ARG
  读写策略  : $RW
  块大小    : $BS
  并发线程  : $NUMJOBS
  队列深度  : $IODEPTH
  文件大小  : $SIZE
  预热时间  : ${RAMP_TIME}s
  测试时长  : ${RUNTIME}s
  结果目录  : $OUTPUT_DIR
================================================

CONF

cleanup() {
    [[ -f "$TEST_FILE" ]] && { echo "[INFO] 清理测试文件: $TEST_FILE"; rm -f "$TEST_FILE"; }
}
trap cleanup EXIT INT TERM

echo "[INFO] 开始预热（${RAMP_TIME}s）+ 正式测试（${RUNTIME}s），请稍候..."

fio \
    --filename="$TEST_FILE" \
    --direct=1 \
    --ioengine=libaio \
    --group_reporting \
    --lockmem=512M \
    --time_based \
    --userspace_reap \
    --randrepeat=0 \
    --norandommap \
    --refill_buffers \
    --rw="$RW" \
    --ramp_time="$RAMP_TIME" \
    --log_avg_msec=1000 \
    --name="$JOB_NAME" \
    --write_lat_log="${OUTPUT_PREFIX}_lat" \
    --write_iops_log="${OUTPUT_PREFIX}_iops" \
    --disable_clat=0 \
    --disable_lat=0 \
    --disable_slat=1 \
    --bs="$BS" \
    --size="$SIZE" \
    --runtime="$RUNTIME" \
    --iodepth="$IODEPTH" \
    --numjobs="$NUMJOBS" \
    --output="$SUMMARY_FILE"

echo ""
echo "[OK] 测试完成，结果文件："
echo "  汇总报告 : $SUMMARY_FILE"
echo "  延迟日志 : ${OUTPUT_PREFIX}_lat.*.log"
echo "  IOPS 日志: ${OUTPUT_PREFIX}_iops.*.log"
echo ""

if [[ -f "$SUMMARY_FILE" ]]; then
    echo "── 关键指标速览 ─────────────────────────────────"
    grep -E "IOPS=|BW=|lat \(|clat|lat_ns" "$SUMMARY_FILE" | head -20 | sed "s/^/  /"
    echo "─────────────────────────────────────────────────"
    echo "  完整结果: cat $SUMMARY_FILE"
    echo ""
fi