#!/usr/bin/env bash
# tcp-bdp-tune.sh —— 按 带宽 × 延迟(BDP) 一键 TCP 调优 (BBR + fq + 缓冲区)
#
# 原理: 要让单条 TCP 在 <带宽> / <RTT延迟> 下跑满, socket 缓冲必须 >= 带宽时延积(BDP)。
#   BDP(bytes) = 带宽(Mbps) × 1e6 × RTT(s) / 8 = Mbps × RTT_ms × 125
#   BBR 拥塞探测期在途数据可达 ~2×BDP, 故缓冲上限取 2×BDP。
#
# 用法:
#   sudo ./tcp-bdp-tune.sh <带宽> <RTT毫秒>      例: sudo ./tcp-bdp-tune.sh 1gbps 100
#   sudo ./tcp-bdp-tune.sh                        不带参数则交互输入
#   ./tcp-bdp-tune.sh 1000 100 --dry-run          只计算+打印, 不改系统 (不需root)
#   sudo ./tcp-bdp-tune.sh --revert               撤销本脚本写入的调优
#
# 带宽单位: 1gbps / 1g / 500mbps / 500m / 纯数字(按Mbps)。RTT 填毫秒。
set -euo pipefail

MARK_BEGIN="# >>> tcp-bdp-tune >>>"
MARK_END="# <<< tcp-bdp-tune <<<"
SYSCTL_FILE="/etc/sysctl.conf"

die(){ echo "错误: $*" >&2; exit 1; }
need_root(){ [ "$(id -u)" = 0 ] || die "需要 root 权限, 请用 sudo 运行"; }

# ---------- 撤销 ----------
if [ "${1:-}" = "--revert" ]; then
  need_root
  if grep -qF "$MARK_BEGIN" "$SYSCTL_FILE" 2>/dev/null; then
    sed -i "\|$MARK_BEGIN|,\|$MARK_END|d" "$SYSCTL_FILE"
    sysctl --system >/dev/null 2>&1 || true
    echo "✅ 已从 $SYSCTL_FILE 移除本脚本的调优块。内核当前运行值需重启后恢复默认(或手动 sysctl 改回)。"
  else
    echo "没找到本脚本写入的调优块, 无需撤销。"
  fi
  exit 0
fi

# ---------- 解析参数 ----------
DRYRUN=0; POS=()
for a in "$@"; do
  case "$a" in
    --dry-run|-n) DRYRUN=1 ;;
    -h|--help) grep -E '^#( |$)' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
    *) POS+=("$a") ;;
  esac
done

parse_bw(){  # -> Mbps
  local s num unit
  s="$(printf '%s' "$1" | tr 'A-Z' 'a-z' | tr -d ' ')"
  num="$(printf '%s' "$s" | grep -oE '^[0-9]+(\.[0-9]+)?')" || die "带宽格式不对: $1"
  unit="${s#"$num"}"
  case "$unit" in
    ''|mbps|mbit|mbit/s|mb|m) awk "BEGIN{print $num}" ;;
    gbps|gbit|gbit/s|gb|g)    awk "BEGIN{print $num*1000}" ;;
    kbps|kbit|kb|k)           awk "BEGIN{print $num/1000}" ;;
    *) die "未知带宽单位 '$unit' (支持 gbps/mbps)";;
  esac
}

BW="${POS[0]:-}"; RTT="${POS[1]:-}"
[ -z "$BW" ]  && read -rp "带宽 (如 1gbps / 500mbps / 1000): " BW
[ -z "$RTT" ] && read -rp "延迟 RTT (毫秒, 如 100): " RTT
MBPS="$(parse_bw "$BW")"
printf '%s' "$RTT" | grep -qE '^[0-9]+(\.[0-9]+)?$' || die "延迟必须是数字(毫秒): $RTT"

# ---------- 计算 BDP / 缓冲 ----------
BDP=$(awk "BEGIN{printf \"%.0f\", $MBPS*$RTT*125}")            # bytes
# 缓冲上限 = 2×BDP, 向上取整到整数 MiB, 下限 4MiB, 上限 512MiB
BUF=$(awk "BEGIN{ b=$BDP*2; mb=b/1048576; mb=(mb==int(mb))?mb:int(mb)+1;
  if(mb<4)mb=4; if(mb>512)mb=512; printf \"%d\", mb*1048576 }")
BUF_MB=$(( BUF / 1048576 ))
BDP_MB=$(awk "BEGIN{printf \"%.2f\", $BDP/1048576}")

echo "──────────────────────────────────────────────"
echo " 输入:  带宽 = ${MBPS} Mbps   延迟 RTT = ${RTT} ms"
echo " BDP  = ${BDP_MB} MiB  (${BDP} bytes)"
echo " 缓冲上限(2×BDP, 取整) = ${BUF_MB} MiB  (rmem_max/wmem_max/tcp_*mem 上限)"
echo "──────────────────────────────────────────────"

BLOCK="$MARK_BEGIN  [带宽${MBPS}Mbps x RTT${RTT}ms, BDP=${BDP_MB}MiB, 缓冲${BUF_MB}MiB]
net.core.default_qdisc=fq
net.ipv4.tcp_congestion_control=bbr
net.core.rmem_max=$BUF
net.core.wmem_max=$BUF
net.core.rmem_default=262144
net.core.wmem_default=262144
net.ipv4.tcp_rmem=4096 131072 $BUF
net.ipv4.tcp_wmem=4096 16384 $BUF
net.ipv4.tcp_slow_start_after_idle=0
net.ipv4.tcp_mtu_probing=1
net.ipv4.tcp_window_scaling=1
net.ipv4.tcp_fastopen=3
net.core.netdev_max_backlog=16384
net.core.somaxconn=8192
$MARK_END"

if [ "$DRYRUN" = 1 ]; then
  echo "[dry-run] 将写入 $SYSCTL_FILE 的内容(未改动系统):"
  echo "$BLOCK"
  exit 0
fi

need_root

# ---------- 环境检查 ----------
if command -v systemd-detect-virt >/dev/null 2>&1 && systemd-detect-virt -c >/dev/null 2>&1; then
  echo "⚠️  检测到容器环境($(systemd-detect-virt))：sysctl 可能只读、BBR 由宿主内核控制, 调优也许不生效。"
fi
modprobe tcp_bbr 2>/dev/null || true
if ! sysctl -n net.ipv4.tcp_available_congestion_control 2>/dev/null | grep -qw bbr; then
  echo "⚠️  当前内核没有 BBR (老内核/容器)。缓冲参数照写, 但拥塞控制不会变成 bbr。"
fi
RAM_MB=$(awk '/MemTotal/{printf "%d",$2/1024}' /proc/meminfo 2>/dev/null || echo 0)
if [ "$RAM_MB" -gt 0 ] && [ "$BUF_MB" -gt $((RAM_MB/8)) ]; then
  echo "⚠️  缓冲上限 ${BUF_MB}MiB 相对内存 ${RAM_MB}MiB 偏大: 单/少量大流没问题, 极高并发下可能吃内存(内核 tcp_mem 会兜底限流)。"
fi

# ---------- 幂等写入 ----------
cp -a "$SYSCTL_FILE" "${SYSCTL_FILE}.bak.$(date +%Y%m%d%H%M%S)" 2>/dev/null || true
[ -f "$SYSCTL_FILE" ] || touch "$SYSCTL_FILE"
# 删旧调优块
sed -i "\|$MARK_BEGIN|,\|$MARK_END|d" "$SYSCTL_FILE" 2>/dev/null || true
# 删散落在 sysctl.conf 里的、本脚本要管的同名键(避免重复/冲突)
for k in net.core.rmem_max net.core.wmem_max net.core.rmem_default net.core.wmem_default \
         net.ipv4.tcp_rmem net.ipv4.tcp_wmem net.core.default_qdisc net.ipv4.tcp_congestion_control \
         net.ipv4.tcp_slow_start_after_idle net.ipv4.tcp_mtu_probing net.ipv4.tcp_window_scaling \
         net.ipv4.tcp_fastopen net.core.netdev_max_backlog net.core.somaxconn; do
  sed -i "/^[[:space:]]*${k//./\\.}[[:space:]]*=/d" "$SYSCTL_FILE" 2>/dev/null || true
done
printf '\n%s\n' "$BLOCK" >> "$SYSCTL_FILE"
# 开机自动加载 bbr 模块
echo tcp_bbr > /etc/modules-load.d/bbr.conf 2>/dev/null || true

# ---------- 应用 ----------
sysctl --system >/dev/null 2>&1 || sysctl -p "$SYSCTL_FILE" >/dev/null 2>&1 || true

# ---------- 验证 ----------
echo "✅ 已应用并写入 $SYSCTL_FILE (重启自动生效, 已备份原文件)。当前生效值:"
for k in net.ipv4.tcp_congestion_control net.core.default_qdisc \
         net.core.rmem_max net.core.wmem_max net.ipv4.tcp_rmem net.ipv4.tcp_wmem \
         net.ipv4.tcp_slow_start_after_idle net.ipv4.tcp_mtu_probing; do
  printf "  %-38s= %s\n" "$k" "$(sysctl -n "$k" 2>/dev/null | tr '\t' ' ')"
done
echo "撤销请运行:  sudo $0 --revert"