499 lines
13 KiB
Go
499 lines
13 KiB
Go
package main
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/json"
|
||
"flag"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"os"
|
||
"time"
|
||
|
||
"plp-test/internal/config"
|
||
"plp-test/internal/model"
|
||
|
||
"bufio"
|
||
"strings"
|
||
|
||
"github.com/google/uuid"
|
||
"github.com/sirupsen/logrus"
|
||
)
|
||
|
||
var (
|
||
configFile string
|
||
logLevel string
|
||
serverAddr string
|
||
testType string
|
||
timeout int
|
||
dataSizeMB int
|
||
blockSize int
|
||
concurrent bool
|
||
recovery bool
|
||
)
|
||
|
||
func init() {
|
||
flag.StringVar(&configFile, "config", "config.yaml", "配置文件路径")
|
||
flag.StringVar(&logLevel, "log-level", "info", "日志级别 (debug, info, warn, error)")
|
||
flag.StringVar(&serverAddr, "server", "", "服务器地址，格式为 host:port")
|
||
flag.StringVar(&testType, "test", "sequential", "测试类型 (sequential, random, mixed, concurrent, power_loss, stability, all)")
|
||
flag.IntVar(&timeout, "timeout", 0, "测试超时时间（秒）")
|
||
flag.IntVar(&dataSizeMB, "data-size", 0, "测试数据大小（MB）")
|
||
flag.IntVar(&blockSize, "block-size", 0, "数据块大小（KB）")
|
||
flag.BoolVar(&concurrent, "concurrent", false, "是否并发执行所有测试")
|
||
flag.BoolVar(&recovery, "recovery", false, "是否请求恢复测试")
|
||
}
|
||
|
||
// Client 客户端
|
||
type Client struct {
|
||
config *config.Config
|
||
logger *logrus.Logger
|
||
httpClient *http.Client
|
||
serverAddr string
|
||
clientID string
|
||
}
|
||
|
||
// NewClient 创建客户端
|
||
func NewClient(cfg *config.Config, logger *logrus.Logger, serverAddr string) *Client {
|
||
if serverAddr == "" {
|
||
serverAddr = cfg.Client.ServerAddr
|
||
}
|
||
|
||
return &Client{
|
||
config: cfg,
|
||
logger: logger,
|
||
httpClient: &http.Client{
|
||
Timeout: time.Duration(cfg.Client.TimeoutSec) * time.Second,
|
||
},
|
||
serverAddr: serverAddr,
|
||
clientID: uuid.New().String(),
|
||
}
|
||
}
|
||
|
||
// RunTest 运行测试
|
||
func (c *Client) RunTest(testType string, dataSizeMB, blockSize int) error {
|
||
c.logger.Infof("运行测试 %s", testType)
|
||
|
||
// 准备请求数据
|
||
req := model.TestRequest{
|
||
TestType: testType,
|
||
DataSizeMB: dataSizeMB,
|
||
BlockSize: blockSize,
|
||
Concurrency: c.config.Client.Concurrency,
|
||
ClientID: c.clientID,
|
||
RequestTime: time.Now(),
|
||
Parameters: make(map[string]string),
|
||
}
|
||
|
||
// 设置默认值
|
||
if req.DataSizeMB == 0 {
|
||
req.DataSizeMB = c.config.Test.DataSizeMB
|
||
}
|
||
if req.BlockSize == 0 {
|
||
req.BlockSize = c.config.Test.BlockSize
|
||
}
|
||
|
||
// 序列化请求数据
|
||
reqData, err := json.Marshal(req)
|
||
if err != nil {
|
||
return fmt.Errorf("序列化请求数据失败: %v", err)
|
||
}
|
||
|
||
// 发送请求
|
||
url := fmt.Sprintf("http://%s/run", c.serverAddr)
|
||
resp, err := c.httpClient.Post(url, "application/json", bytes.NewBuffer(reqData))
|
||
if err != nil {
|
||
return fmt.Errorf("发送请求失败: %v", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
// 检查响应状态
|
||
if resp.StatusCode != http.StatusAccepted {
|
||
return fmt.Errorf("服务器返回错误状态码: %d", resp.StatusCode)
|
||
}
|
||
|
||
// 解析响应
|
||
var testResp model.TestResponse
|
||
if err := json.NewDecoder(resp.Body).Decode(&testResp); err != nil {
|
||
return fmt.Errorf("解析响应失败: %v", err)
|
||
}
|
||
|
||
c.logger.Infof("测试请求已接受，RequestID: %s", testResp.RequestID)
|
||
|
||
// 监控测试状态
|
||
return c.MonitorTestStatus(testResp.RequestID)
|
||
}
|
||
|
||
// MonitorTestStatus 监控测试状态
|
||
func (c *Client) MonitorTestStatus(testID string) error {
|
||
c.logger.Infof("监控测试状态: %s", testID)
|
||
|
||
// 使用流式API监控进度
|
||
errCh := make(chan error, 1)
|
||
doneCh := make(chan struct{})
|
||
|
||
// 启动流式监控
|
||
go func() {
|
||
defer close(doneCh)
|
||
|
||
url := fmt.Sprintf("http://%s/stream?test_id=%s&client_id=%s", c.serverAddr, testID, c.clientID)
|
||
|
||
// 创建一个没有超时的HTTP客户端专门用于流式连接
|
||
noTimeoutClient := &http.Client{
|
||
Timeout: 0, // 0表示无超时
|
||
}
|
||
|
||
req, err := http.NewRequest("GET", url, nil)
|
||
if err != nil {
|
||
errCh <- fmt.Errorf("创建请求失败: %v", err)
|
||
return
|
||
}
|
||
|
||
resp, err := noTimeoutClient.Do(req)
|
||
if err != nil {
|
||
errCh <- fmt.Errorf("发送流请求失败: %v", err)
|
||
return
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode != http.StatusOK {
|
||
errCh <- fmt.Errorf("服务器返回错误状态码: %d", resp.StatusCode)
|
||
return
|
||
}
|
||
|
||
// 创建事件扫描器
|
||
scanner := bufio.NewScanner(resp.Body)
|
||
|
||
// 处理Server-Sent Events
|
||
for scanner.Scan() {
|
||
line := scanner.Text()
|
||
|
||
// 跳过空行
|
||
if line == "" {
|
||
continue
|
||
}
|
||
|
||
// 处理数据行
|
||
if strings.HasPrefix(line, "data: ") {
|
||
data := line[6:] // 去掉 "data: " 前缀
|
||
|
||
// 解析更新数据
|
||
var update model.StreamUpdate
|
||
if err := json.Unmarshal([]byte(data), &update); err != nil {
|
||
c.logger.Warnf("解析更新数据失败: %v", err)
|
||
continue
|
||
}
|
||
|
||
// 处理不同类型的更新
|
||
switch update.Type {
|
||
case "status":
|
||
c.logger.Infof("测试状态: 进度: %.2f%%, 阶段: %s",
|
||
update.Progress, update.CurrentPhase)
|
||
case "error":
|
||
c.logger.Errorf("测试错误: %s", update.Message)
|
||
case "completion":
|
||
c.logger.Infof("测试完成: %s", update.Message)
|
||
return
|
||
case "integrity":
|
||
if info, ok := update.Data.(map[string]interface{}); ok {
|
||
c.logger.Infof("数据完整性: 可用块: %.0f, 损坏块: %.0f, 丢失块: %.0f, 数据丢失: %.2f MB",
|
||
info["available_blocks"], info["corrupted_blocks"], info["missing_blocks"], info["data_loss_mb"])
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
if err := scanner.Err(); err != nil {
|
||
// 如果是客户端主动断开连接，不报错
|
||
if !strings.Contains(err.Error(), "use of closed network connection") {
|
||
errCh <- fmt.Errorf("读取流数据失败: %v", err)
|
||
}
|
||
}
|
||
}()
|
||
|
||
// 等待流结束或出错
|
||
select {
|
||
case err := <-errCh:
|
||
// 如果流模式失败，回退到轮询模式
|
||
c.logger.Warnf("流式监控失败: %v, 回退到轮询模式", err)
|
||
return c.pollTestStatus(testID)
|
||
case <-doneCh:
|
||
c.logger.Info("流式监控完成")
|
||
return nil
|
||
}
|
||
}
|
||
|
||
// pollTestStatus 使用轮询方式监控测试状态
|
||
func (c *Client) pollTestStatus(testID string) error {
|
||
c.logger.Infof("使用轮询监控测试状态: %s", testID)
|
||
|
||
for {
|
||
// 获取测试状态
|
||
url := fmt.Sprintf("http://%s/status?test_id=%s", c.serverAddr, testID)
|
||
resp, err := c.httpClient.Get(url)
|
||
if err != nil {
|
||
c.logger.Warnf("获取测试状态失败: %v", err)
|
||
time.Sleep(2 * time.Second)
|
||
continue
|
||
}
|
||
|
||
// 检查响应状态
|
||
if resp.StatusCode == http.StatusNotFound {
|
||
c.logger.Infof("测试 %s 已完成", testID)
|
||
resp.Body.Close()
|
||
break
|
||
}
|
||
|
||
if resp.StatusCode != http.StatusOK {
|
||
resp.Body.Close()
|
||
c.logger.Warnf("服务器返回错误状态码: %d", resp.StatusCode)
|
||
time.Sleep(2 * time.Second)
|
||
continue
|
||
}
|
||
|
||
// 解析响应
|
||
var status model.TestStatus
|
||
if err := json.NewDecoder(resp.Body).Decode(&status); err != nil {
|
||
resp.Body.Close()
|
||
c.logger.Warnf("解析响应失败: %v", err)
|
||
time.Sleep(2 * time.Second)
|
||
continue
|
||
}
|
||
resp.Body.Close()
|
||
|
||
// 显示测试状态
|
||
c.logger.Infof("测试状态: %s, 进度: %.2f%%, 阶段: %s",
|
||
status.Status, status.Progress, status.CurrentPhase)
|
||
|
||
// 检查测试是否结束
|
||
if status.Status == "completed" || status.Status == "failed" || status.Status == "aborted" {
|
||
c.logger.Infof("测试 %s %s", testID, status.Status)
|
||
break
|
||
}
|
||
|
||
// 等待一段时间再次检查
|
||
time.Sleep(1 * time.Second)
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// CheckServerHealth 检查服务器健康状态
|
||
func (c *Client) CheckServerHealth() error {
|
||
c.logger.Info("检查服务器健康状态")
|
||
|
||
url := fmt.Sprintf("http://%s/health", c.serverAddr)
|
||
resp, err := c.httpClient.Get(url)
|
||
if err != nil {
|
||
return fmt.Errorf("连接服务器失败: %v", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode != http.StatusOK {
|
||
return fmt.Errorf("服务器返回错误状态码: %d", resp.StatusCode)
|
||
}
|
||
|
||
var health model.HealthStatus
|
||
if err := json.NewDecoder(resp.Body).Decode(&health); err != nil {
|
||
return fmt.Errorf("解析响应失败: %v", err)
|
||
}
|
||
|
||
c.logger.Infof("服务器状态: %s, 消息: %s", health.Status, health.Message)
|
||
return nil
|
||
}
|
||
|
||
// RunAllTests 运行所有测试
|
||
func (c *Client) RunAllTests(concurrent bool) error {
|
||
c.logger.Info("运行所有测试")
|
||
|
||
tests := c.config.Test.EnabledTests
|
||
|
||
if concurrent {
|
||
c.logger.Info("并发执行所有测试")
|
||
errCh := make(chan error, len(tests))
|
||
|
||
for _, test := range tests {
|
||
go func(t string) {
|
||
errCh <- c.RunTest(t, dataSizeMB, blockSize)
|
||
}(test)
|
||
}
|
||
|
||
// 等待所有测试完成
|
||
for range tests {
|
||
if err := <-errCh; err != nil {
|
||
c.logger.Errorf("测试失败: %v", err)
|
||
}
|
||
}
|
||
} else {
|
||
c.logger.Info("顺序执行所有测试")
|
||
for _, test := range tests {
|
||
if err := c.RunTest(test, dataSizeMB, blockSize); err != nil {
|
||
c.logger.Errorf("测试 %s 失败: %v", test, err)
|
||
}
|
||
}
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// CheckDataIntegrity 检查断电后的数据完整性
|
||
func (c *Client) CheckDataIntegrity(testID string) error {
|
||
c.logger.Info("检查数据完整性")
|
||
|
||
url := fmt.Sprintf("http://%s/integrity?test_id=%s", c.serverAddr, testID)
|
||
resp, err := c.httpClient.Get(url)
|
||
if err != nil {
|
||
return fmt.Errorf("获取数据完整性信息失败: %v", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode != http.StatusOK {
|
||
return fmt.Errorf("服务器返回错误状态码: %d", resp.StatusCode)
|
||
}
|
||
|
||
var integrity model.IntegrityInfo
|
||
if err := json.NewDecoder(resp.Body).Decode(&integrity); err != nil {
|
||
return fmt.Errorf("解析响应失败: %v", err)
|
||
}
|
||
|
||
c.logger.Infof("数据完整性检查结果:")
|
||
c.logger.Infof(" 测试ID: %s", integrity.TestID)
|
||
c.logger.Infof(" 检查时间: %s", integrity.CheckTime.Format("2006-01-02 15:04:05"))
|
||
c.logger.Infof(" 总块数: %d", integrity.TotalBlocks)
|
||
c.logger.Infof(" 可用块数: %d (%.2f%%)", integrity.AvailableBlocks, float64(integrity.AvailableBlocks)/float64(integrity.TotalBlocks)*100)
|
||
c.logger.Infof(" 损坏块数: %d (%.2f%%)", integrity.CorruptedBlocks, float64(integrity.CorruptedBlocks)/float64(integrity.TotalBlocks)*100)
|
||
c.logger.Infof(" 丢失块数: %d (%.2f%%)", integrity.MissingBlocks, float64(integrity.MissingBlocks)/float64(integrity.TotalBlocks)*100)
|
||
c.logger.Infof(" 数据丢失: %.2f MB", integrity.DataLossMB)
|
||
c.logger.Infof(" 恢复成功: %v", integrity.RecoverySuccess)
|
||
|
||
return nil
|
||
}
|
||
|
||
// RequestRecoveryTest 请求执行恢复测试
|
||
func (c *Client) RequestRecoveryTest(testType string) error {
|
||
c.logger.Infof("请求恢复测试: %s", testType)
|
||
|
||
// 准备请求数据
|
||
req := struct {
|
||
TestType string `json:"test_type"`
|
||
TestDir string `json:"test_dir,omitempty"`
|
||
}{
|
||
TestType: testType,
|
||
}
|
||
|
||
// 序列化请求数据
|
||
reqData, err := json.Marshal(req)
|
||
if err != nil {
|
||
return fmt.Errorf("序列化请求数据失败: %v", err)
|
||
}
|
||
|
||
// 发送请求
|
||
url := fmt.Sprintf("http://%s/recovery", c.serverAddr)
|
||
resp, err := c.httpClient.Post(url, "application/json", bytes.NewBuffer(reqData))
|
||
if err != nil {
|
||
return fmt.Errorf("发送请求失败: %v", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
// 检查响应状态
|
||
if resp.StatusCode != http.StatusAccepted {
|
||
return fmt.Errorf("服务器返回错误状态码: %d", resp.StatusCode)
|
||
}
|
||
|
||
// 解析响应
|
||
var testResp model.TestResponse
|
||
if err := json.NewDecoder(resp.Body).Decode(&testResp); err != nil {
|
||
return fmt.Errorf("解析响应失败: %v", err)
|
||
}
|
||
|
||
c.logger.Infof("恢复测试请求已接受，RequestID: %s", testResp.RequestID)
|
||
|
||
// 监控测试状态
|
||
return c.MonitorTestStatus(testResp.RequestID)
|
||
}
|
||
|
||
func main() {
|
||
flag.Parse()
|
||
|
||
// 初始化日志级别
|
||
var level logrus.Level
|
||
switch logLevel {
|
||
case "debug":
|
||
level = logrus.DebugLevel
|
||
case "info":
|
||
level = logrus.InfoLevel
|
||
case "warn":
|
||
level = logrus.WarnLevel
|
||
case "error":
|
||
level = logrus.ErrorLevel
|
||
default:
|
||
level = logrus.InfoLevel
|
||
}
|
||
|
||
// 配置日志
|
||
logger := logrus.New()
|
||
logger.SetLevel(level)
|
||
logger.SetFormatter(&logrus.TextFormatter{
|
||
FullTimestamp: true,
|
||
TimestampFormat: "2006-01-02 15:04:05",
|
||
})
|
||
|
||
// 添加日志文件输出
|
||
logFile, err := os.OpenFile("client.log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
|
||
if err == nil {
|
||
multiWriter := io.MultiWriter(os.Stdout, logFile)
|
||
logger.SetOutput(multiWriter)
|
||
} else {
|
||
logger.Warnf("无法打开日志文件: %v", err)
|
||
}
|
||
|
||
// 加载配置
|
||
cfg, err := config.Load(configFile)
|
||
if err != nil {
|
||
logger.Fatalf("加载配置失败: %v", err)
|
||
}
|
||
|
||
// 如果指定了timeout，覆盖配置
|
||
if timeout > 0 {
|
||
cfg.Client.TimeoutSec = timeout
|
||
}
|
||
|
||
// 创建客户端
|
||
client := NewClient(cfg, logger, serverAddr)
|
||
|
||
// 检查服务器健康状态
|
||
if err := client.CheckServerHealth(); err != nil {
|
||
logger.Fatalf("服务器健康检查失败: %v", err)
|
||
}
|
||
|
||
// 处理恢复测试
|
||
if recovery {
|
||
if testType == "" {
|
||
testType = "power_loss"
|
||
}
|
||
logger.Infof("开始执行恢复测试: %s", testType)
|
||
if err := client.RequestRecoveryTest(testType); err != nil {
|
||
logger.Fatalf("恢复测试失败: %v", err)
|
||
}
|
||
logger.Info("恢复测试完成")
|
||
return
|
||
}
|
||
|
||
// 执行指定的测试
|
||
if testType != "" {
|
||
if testType == "all" {
|
||
if err := client.RunAllTests(concurrent); err != nil {
|
||
logger.Fatalf("执行所有测试失败: %v", err)
|
||
}
|
||
} else {
|
||
if err := client.RunTest(testType, dataSizeMB, blockSize); err != nil {
|
||
logger.Fatalf("执行测试 %s 失败: %v", testType, err)
|
||
}
|
||
}
|
||
} else {
|
||
logger.Fatal("未指定测试类型，请使用 -test 参数")
|
||
}
|
||
|
||
logger.Info("所有测试完成")
|
||
}
|