Golang压测ws

发布时间 2023-12-25 10:51:29作者: 朝阳1
package main

import (
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

func main() {
	// 设置要压测的WebSocket服务地址
	url := "ws://192.168.252.128:8080/"

	// 设置并发请求数量
	concurrency := 100

	// 设置每个连接的发送消息数量
	messageCount := 10

	// 等待组,用于等待所有goroutine完成
	var wg sync.WaitGroup

	// 开始时间
	startTime := time.Now()

	// 根据并发请求数量启动goroutine
	for i := 0; i < concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()

			// 建立连接
			conn, _, err := websocket.DefaultDialer.Dial(url, nil)
			if err != nil {
				log.Fatal("无法建立WebSocket连接:", err)
			}
			defer conn.Close()

			// 发送指定数量的消息
			for j := 0; j < messageCount; j++ {
				message := fmt.Sprintf("这是第 %d 个消息", j)
				err = conn.WriteMessage(websocket.TextMessage, []byte(message))
				if err != nil {
					log.Println("发送消息失败:", err)
					return
				}
				log.Println("发送消息:", message)

				// 在发送每个消息之后,可以在这里等待一段时间以模拟实际场景的负载
				//time.Sleep(time.Second)
			}

			// 读取服务器响应
			_, message, err := conn.ReadMessage()
			if err != nil {
				log.Println("读取服务器响应失败:", err)
				return
			}
			log.Println("服务器响应:", string(message))
		}()
	}

	// 等待所有goroutine完成
	wg.Wait()

	// 计算执行时间
	executionTime := time.Since(startTime)

	// 打印统计信息
	log.Printf("压测完成,共发送 %d 个消息,每个连接 %d 个消息\n", concurrency*messageCount, messageCount)
	log.Printf("总执行时间: %.2f 秒\n", executionTime.Seconds())
	log.Printf("每秒请求数: %.2f\n", float64(concurrency*messageCount)/executionTime.Seconds())
}