#!/bin/bash

# DNS性能优化验证脚本
SERVER="127.0.0.1"
DOMAIN="example.com"
THREADS=50
QUERIES=5000

echo "========================================="
echo "DNS服务器多线程优化验证测试"
echo "========================================="
echo "服务器: $SERVER"
echo "域名: $DOMAIN"
echo "线程数: $THREADS"
echo "总查询数: $QUERIES"
echo "========================================="

echo ""
echo "启动DNS服务器..."
./dns-server > /dev/null 2>&1 &
DNS_PID=$!
echo "DNS服务器PID: $DNS_PID"

sleep 5

echo ""
echo "========================================="
echo "开始压力测试前系统状态："
echo "========================================="
echo "CPU使用率:"
top -b -n 1 | grep "Cpu(s)" | awk '{print "  " $2 "% 用户, " $4 "% 系统, " $8 "% 空闲"}'
echo ""
echo "内存使用情况:"
free -h | grep -E "Mem|Swap" | awk '{print "  " $1 ": " $3 "/" $2 " (" $3/$2*100 "% 使用)"}'
echo ""
echo "DNS服务器进程资源使用:"
ps -p $DNS_PID -o %cpu,%mem,cmd --no-headers | awk '{print "  CPU: " $1 "%, 内存: " $2 "%"}'

echo ""
echo "========================================="
echo "开始压力测试..."
echo "========================================="

start_time=$(date +%s.%N)

pids=()
for ((i=1; i<=THREADS; i++)); do
  for ((j=1; j<=$((QUERIES/THREADS)); j++)); do
    dig @$SERVER $DOMAIN A +short > /dev/null 2>&1 &
    pids+=($!)
  done
  if (( $i % 10 == 0 )); then
    echo "已启动 $i/$THREADS 个线程..."
    sleep 0.2
  fi
done

echo "所有线程已启动，等待完成..."

for pid in "${pids[@]}"; do
  wait $pid 2>/dev/null
done

end_time=$(date +%s.%N)
elapsed=$(echo "$end_time - $start_time" | bc)
qps=$(echo "$QUERIES / $elapsed" | bc)

echo ""
echo "========================================="
echo "压力测试后系统状态："
echo "========================================="
echo "CPU使用率:"
top -b -n 1 | grep "Cpu(s)" | awk '{print "  " $2 "% 用户, " $4 "% 系统, " $8 "% 空闲"}'
echo ""
echo "内存使用情况:"
free -h | grep -E "Mem|Swap" | awk '{print "  " $1 ": " $3 "/" $2 " (" $3/$2*100 "% 使用)"}'
echo ""
echo "DNS服务器进程资源使用:"
ps -p $DNS_PID -o %cpu,%mem,cmd --no-headers 2>/dev/null | awk '{print "  CPU: " $1 "%, 内存: " $2 "%"}'

echo ""
echo "========================================="
echo "测试结果："
echo "========================================="
echo "总查询数: $QUERIES"
echo "执行时间: $elapsed 秒"
echo "QPS: $qps 次/秒"
echo "========================================="

echo ""
echo "停止DNS服务器..."
kill $DNS_PID 2>/dev/null
sleep 2

echo ""
echo "测试完成！"