Files
TencentDB-Agent-Memory/scripts/stress-test-v2.sh
T

326 lines
12 KiB
Bash
Raw Normal View History

2026-05-29 17:33:12 +08:00
#!/usr/bin/env bash
# ============================================================
# V2 API 全接口压测脚本
# 测试所有12个 v2 端点的可用性、并发处理、Redis 状态正确性
# ============================================================
set -euo pipefail
HOST="${1:-http://localhost:8420}"
API_KEY="${2:-b9WmawnJFpb9vn0XWKDKSxF5Eaf5SeXdIHaRpShmSbgg}"
SERVICE_ID="${3:-mem-j4wjesud}"
CONCURRENCY="${4:-10}"
ROUNDS="${5:-5}"
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'
PASS=0
FAIL=0
ERRORS=""
log_ok() { echo -e "${GREEN}[PASS]${NC} $1"; PASS=$((PASS+1)); }
log_fail() { echo -e "${RED}[FAIL]${NC} $1"; FAIL=$((FAIL+1)); ERRORS="${ERRORS}\n - $1"; }
log_info() { echo -e "${YELLOW}[INFO]${NC} $1"; }
call_api() {
local endpoint="$1"
local body="$2"
local desc="$3"
local resp
resp=$(curl -s -w "\n%{http_code}" -X POST "${HOST}${endpoint}" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer ${API_KEY}" \
-H "x-tdai-service-id: ${SERVICE_ID}" \
-d "${body}" 2>/dev/null)
local http_code
http_code=$(echo "$resp" | tail -1)
local body_resp
body_resp=$(echo "$resp" | sed '$d')
local code
code=$(echo "$body_resp" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('code',''))" 2>/dev/null || echo "parse_error")
if [[ "$http_code" == "200" && "$code" == "0" ]]; then
log_ok "$desc (HTTP=$http_code, code=$code)"
echo "$body_resp"
return 0
else
log_fail "$desc (HTTP=$http_code, code=$code, body=${body_resp:0:200})"
return 1
fi
}
# ============================================================
echo "=========================================="
echo " V2 API 全接口压测"
echo " Host: $HOST"
echo " ServiceID: $SERVICE_ID"
echo " Concurrency: $CONCURRENCY"
echo " Rounds: $ROUNDS"
echo "=========================================="
# 0. Health check
log_info "Step 0: Health check"
HEALTH=$(curl -s "${HOST}/health" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d['status'])" 2>/dev/null)
if [[ "$HEALTH" == "ok" ]]; then
log_ok "Health check"
else
log_fail "Health check (status=$HEALTH)"
exit 1
fi
# ============================================================
# 1. Conversation Add (L0 写入) - 并发压测
# ============================================================
log_info "Step 1: POST /v2/conversation/add — $CONCURRENCY concurrent × $ROUNDS rounds"
TS=$(date +%s)
for round in $(seq 1 $ROUNDS); do
pids=()
for i in $(seq 1 $CONCURRENCY); do
SESSION="stress-${TS}-r${round}-c${i}"
BODY=$(cat <<EOF
{
"instance_id": "${SERVICE_ID}",
"session_id": "${SESSION}",
"messages": [
{"role": "user", "content": "压测消息 round=${round} thread=${i} 我喜欢在周末去公园跑步"},
{"role": "assistant", "content": "跑步是很好的运动习惯!你通常跑多久?"},
{"role": "user", "content": "一般跑30分钟左右,然后做拉伸"},
{"role": "assistant", "content": "30分钟加拉伸是很健康的安排"}
]
}
EOF
)
curl -s -o /dev/null -w "" -X POST "${HOST}/v2/conversation/add" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer ${API_KEY}" \
-H "x-tdai-service-id: ${SERVICE_ID}" \
-d "${BODY}" &
pids+=($!)
done
# Wait for all concurrent requests
fail_count=0
for pid in "${pids[@]}"; do
wait "$pid" || fail_count=$((fail_count+1))
done
if [[ $fail_count -eq 0 ]]; then
log_ok "Round $round: $CONCURRENCY concurrent conversation/add OK"
else
log_fail "Round $round: $fail_count/$CONCURRENCY conversation/add failed"
fi
done
# ============================================================
# 2. Conversation Query (L0 查询)
# ============================================================
log_info "Step 2: POST /v2/conversation/query"
SESSION_Q="stress-${TS}-r1-c1"
call_api "/v2/conversation/query" \
"{\"instance_id\":\"${SERVICE_ID}\",\"session_id\":\"${SESSION_Q}\",\"limit\":10}" \
"conversation/query session=${SESSION_Q}" > /dev/null || true
# ============================================================
# 3. Conversation Search (L0 搜索)
# ============================================================
log_info "Step 3: POST /v2/conversation/search"
call_api "/v2/conversation/search" \
"{\"instance_id\":\"${SERVICE_ID}\",\"query\":\"跑步\",\"limit\":5}" \
"conversation/search query=跑步" > /dev/null || true
# ============================================================
# 4. Conversation Delete (L0 删除)
# ============================================================
log_info "Step 4: POST /v2/conversation/delete"
call_api "/v2/conversation/delete" \
"{\"instance_id\":\"${SERVICE_ID}\",\"session_id\":\"stress-${TS}-r${ROUNDS}-c${CONCURRENCY}\",\"ids\":[\"nonexistent-id\"]}" \
"conversation/delete (nonexistent OK)" > /dev/null || true
# ============================================================
# 5. Atomic Add (L1 写入)
# ============================================================
log_info "Step 5: POST /v2/atomic/add — concurrent"
pids=()
for i in $(seq 1 $CONCURRENCY); do
BODY=$(cat <<EOF
{
"instance_id": "${SERVICE_ID}",
"session_id": "stress-atomic-${TS}",
"memories": [
{"text": "用户喜欢在周末跑步 (thread=$i)", "type": "episodic", "priority": 0.8},
{"text": "用户每次跑30分钟 (thread=$i)", "type": "episodic", "priority": 0.6}
]
}
EOF
)
curl -s -o /dev/null -w "" -X POST "${HOST}/v2/atomic/add" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer ${API_KEY}" \
-H "x-tdai-service-id: ${SERVICE_ID}" \
-d "${BODY}" &
pids+=($!)
done
fail_count=0
for pid in "${pids[@]}"; do
wait "$pid" || fail_count=$((fail_count+1))
done
if [[ $fail_count -eq 0 ]]; then
log_ok "atomic/add: $CONCURRENCY concurrent OK"
else
log_fail "atomic/add: $fail_count/$CONCURRENCY failed"
fi
# ============================================================
# 6. Atomic Query (L1 查询)
# ============================================================
log_info "Step 6: POST /v2/atomic/query"
call_api "/v2/atomic/query" \
"{\"instance_id\":\"${SERVICE_ID}\",\"session_id\":\"stress-atomic-${TS}\",\"limit\":20}" \
"atomic/query" > /dev/null || true
# ============================================================
# 7. Atomic Search (L1 搜索)
# ============================================================
log_info "Step 7: POST /v2/atomic/search"
call_api "/v2/atomic/search" \
"{\"instance_id\":\"${SERVICE_ID}\",\"query\":\"跑步\",\"limit\":5}" \
"atomic/search query=跑步" > /dev/null || true
# ============================================================
# 8. Atomic Delete (L1 删除)
# ============================================================
log_info "Step 8: POST /v2/atomic/delete"
call_api "/v2/atomic/delete" \
"{\"instance_id\":\"${SERVICE_ID}\",\"ids\":[\"nonexistent-l1-id\"]}" \
"atomic/delete (nonexistent OK)" > /dev/null || true
# ============================================================
# 9. Scenario List (L2 列表)
# ============================================================
log_info "Step 9: POST /v2/scenario/ls"
call_api "/v2/scenario/ls" \
"{\"instance_id\":\"${SERVICE_ID}\"}" \
"scenario/ls" > /dev/null || true
# ============================================================
# 10. Scenario Read (L2 读取)
# ============================================================
log_info "Step 10: POST /v2/scenario/read"
call_api "/v2/scenario/read" \
"{\"instance_id\":\"${SERVICE_ID}\",\"path\":\"failover-test.md\"}" \
"scenario/read path=failover-test.md" > /dev/null || true
# ============================================================
# 11. Scenario Write (L2 写入)
# ============================================================
log_info "Step 11: POST /v2/scenario/write"
call_api "/v2/scenario/write" \
"{\"instance_id\":\"${SERVICE_ID}\",\"path\":\"stress-test-${TS}.md\",\"content\":\"# 压测场景\\n\\n这是压测写入的场景文件。\\n\\n## 用户偏好\\n- 周末跑步\"}" \
"scenario/write" > /dev/null || true
# ============================================================
# 12. Scenario Delete (L2 删除)
# ============================================================
log_info "Step 12: POST /v2/scenario/rm"
call_api "/v2/scenario/rm" \
"{\"instance_id\":\"${SERVICE_ID}\",\"path\":\"stress-test-${TS}.md\"}" \
"scenario/rm (cleanup)" > /dev/null || true
# ============================================================
# 13. Persona Read (L3 读取)
# ============================================================
log_info "Step 13: POST /v2/persona/read"
call_api "/v2/persona/read" \
"{\"instance_id\":\"${SERVICE_ID}\"}" \
"persona/read" > /dev/null || true
# ============================================================
# 14. Persona Write (L3 写入)
# ============================================================
log_info "Step 14: POST /v2/persona/write"
call_api "/v2/persona/write" \
"{\"instance_id\":\"${SERVICE_ID}\",\"content\":\"# 用户画像\\n\\n## 基本信息\\n- 喜欢跑步\\n- 周末锻炼\\n\"}" \
"persona/write" > /dev/null || true
# ============================================================
# 15. 高并发混合读写压测
# ============================================================
log_info "Step 15: Mixed concurrent read/write stress ($((CONCURRENCY*2)) requests)"
pids=()
for i in $(seq 1 $CONCURRENCY); do
# Write
curl -s -o /dev/null -X POST "${HOST}/v2/conversation/add" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer ${API_KEY}" \
-H "x-tdai-service-id: ${SERVICE_ID}" \
-d "{\"instance_id\":\"${SERVICE_ID}\",\"session_id\":\"mix-${TS}-${i}\",\"messages\":[{\"role\":\"user\",\"content\":\"混合压测消息 $i\"}]}" &
pids+=($!)
# Read
curl -s -o /dev/null -X POST "${HOST}/v2/conversation/query" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer ${API_KEY}" \
-H "x-tdai-service-id: ${SERVICE_ID}" \
-d "{\"instance_id\":\"${SERVICE_ID}\",\"session_id\":\"mix-${TS}-${i}\",\"limit\":5}" &
pids+=($!)
done
fail_count=0
for pid in "${pids[@]}"; do
wait "$pid" || fail_count=$((fail_count+1))
done
if [[ $fail_count -eq 0 ]]; then
log_ok "Mixed stress: $((CONCURRENCY*2)) concurrent OK"
else
log_fail "Mixed stress: $fail_count/$((CONCURRENCY*2)) failed"
fi
# ============================================================
# 16. Redis 状态校验
# ============================================================
log_info "Step 16: Health check + Redis state validation"
HEALTH_RESP=$(curl -s "${HOST}/health")
WORKER_TASKS=$(echo "$HEALTH_RESP" | python3 -c "import sys,json; d=json.load(sys.stdin); s=d['services']; print(f\"consumed={s['pipelineWorker']['tasksConsumed']} completed={s['pipelineWorker']['tasksCompleted']} failed={s['pipelineWorker']['tasksFailed']} deadLetter={s['pipelineWorker']['tasksDeadLettered']}\")" 2>/dev/null)
SCANNER=$(echo "$HEALTH_RESP" | python3 -c "import sys,json; d=json.load(sys.stdin); s=d['services']['timerScanner']; print(f\"scans={s['scansCompleted']} enqueued={s['tasksEnqueued']} errors={s['scanErrors']}\")" 2>/dev/null)
STATE=$(echo "$HEALTH_RESP" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d['services']['stateBackend'])" 2>/dev/null)
echo " Worker: $WORKER_TASKS"
echo " Scanner: $SCANNER"
echo " Redis: $STATE"
if [[ "$STATE" == "connected" ]]; then
log_ok "Redis state: connected"
else
log_fail "Redis state: $STATE"
fi
TASK_FAILED=$(echo "$HEALTH_RESP" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d['services']['pipelineWorker']['tasksFailed'])" 2>/dev/null)
if [[ "$TASK_FAILED" == "0" ]]; then
log_ok "No pipeline task failures"
else
log_fail "Pipeline task failures: $TASK_FAILED"
fi
# ============================================================
# Summary
# ============================================================
echo ""
echo "=========================================="
echo " 压测结果汇总"
echo "=========================================="
echo -e " ${GREEN}PASS: $PASS${NC}"
echo -e " ${RED}FAIL: $FAIL${NC}"
if [[ $FAIL -gt 0 ]]; then
echo -e " 失败详情:${ERRORS}"
fi
echo "=========================================="
TOTAL_MSGS=$((CONCURRENCY * ROUNDS * 4 + CONCURRENCY * 1))
echo " 总写入消息数: ~$TOTAL_MSGS (L0) + $((CONCURRENCY*2)) (L1)"
echo "=========================================="
exit $FAIL