【go标准库系列】go语言实现简单的压力测试工具

业务背景

golang实现对某个接口发送http请求，用于测试java应用程序的性能。

源码和使用方法

工具用法如下：

./main 目标地址 目标文件夹/postdata.txt 119 目标文件夹/uidsFile.txt

代码如下

package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

func main() {
	url := os.Args[1]
	filePath := os.Args[2]

	// 定义要发起的请求数量
	numRequests, _ := strconv.Atoi(os.Args[3])
	uidsFile, _ := ioutil.ReadFile(os.Args[4])
	fmt.Printf("Request file is  %s ", string(uidsFile))
	uids := strings.Split(string(uidsFile), ",")
	// 创建等待组，用于等待所有请求完成
	var wg sync.WaitGroup
	wg.Add(numRequests)

	// 定义通道，用于协调goroutine之间的通信
	ch := make(chan struct{}, numRequests)


	jsonBytes, _ := ioutil.ReadFile(filePath)

	// 发起并发请求
	for _, uid := range uids {
		newUid, _ := strconv.ParseInt(uid, 10, 64)
		go func(newUid int64) {
			// 将goroutine添加到等待组，以便在结束时通知等待组
			defer wg.Done()

			newString := strings.Replace(string(jsonBytes), "uiddata", strconv.FormatInt(newUid, 10), -1)

			// 构造HTTP请求
			req, err := http.NewRequest("POST", url, bytes.NewBuffer([]byte(newString)))
			if err != nil {
				fmt.Printf("Error creating request: %v\n", err)
				return
			}
			req.Header.Set("Content-Type", "application/json")

			// 发送HTTP请求
			start := time.Now()
			client := &http.Client{}
			resp, err := client.Do(req)
			if err != nil {
				fmt.Printf("Error sending request: %v\n", err)
				return
			}
			body, err := ioutil.ReadAll(resp.Body)
			fmt.Println("response", string(body))
			defer resp.Body.Close()

			// 打印请求耗时
			fmt.Printf("Request %d took %v\n", newUid, time.Since(start))

			// 向通道发送信号，表示该请求已完成
			ch <- struct{}{}
		}(newUid)
	}

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

	// 关闭通道
	close(ch)

	// 等待所有goroutine接收通道信号并退出
	for range ch {
	}

	fmt.Println("All requests completed.")
}