基于go语言的http压力测试工具-代码篇

序言

简单介绍开发背景：基于go语言的http压力测试工具-简单介绍_SteveGao的博客-CSDN博客

花了2个礼拜看了下go语言，第1周主要看了看基础语法和协程相关的知识，第2周了解了下go-zero框架和gRPC，并了解 proto buffer，这周为了加强印象，同时也为了尝鲜go的并发能力，想着就模仿webbench或者apache ab 写一个相似的压测工具，自己瞎折腾，我觉得只有写才知道哪些地方没有掌握，实践出真知吧。仅仅记录自己学习成果，防止后面忘记找不到了。

代码目录：

-config 目录保存压测配置

-log 保存请求日志

-out ide打包输出目录

-process 源码存放目录

-- config.go 解析配置文件 config/config.ini

-- futil.go 文件读取操作

-- httprequest.go http 请求 post or get

-- logger.go 日志输出到文件配置目录

-- parse.go 命令行参数解析

-- request.go 协程调用

-- response.go 接收并处理返回

-- tool.go    main.go 中一些封装函数

--main.go 主函数

config/config.ini

配置文件的简单解释：

配置文件参数部分来自ab，部分来自curl，都比较常见，断言想法来自jmeter。

[process]
;;请求地址
url=http://domain/user/login/login
;;请求方式
x=post
;;开启协程数量
c=3
;;每个协程请求次数
n=1
;;header头
h='Content-Type:application/json'
;;请求参数
d={"mobile":"6712331123", "code" : 888888, "appId":1}
;;累计请求时间
;t=10
;;断言成功的配置
assertion="code":0

main.go 

在主函数中，主要做了如下几件事情：

1、接收命令行的参数 或者 来自配置文件的参数

2、使用sync.WaitGroup 来确保每轮的请求都等待完成，因为如果主协程退出后，其他子协程也跟着退出

3、使用for循环来重复执行 请求，直到 执行时间 等于 -t 设置的时间

4、调用tool.go 打印执行结果

package main

import (
	"fmt"
	"log"
	"os"
	"processtest/process"
	"runtime"
	"sync"
	"time"
)

var lf process.LogFile

func init() {
	lf = process.LogFile{
		"log.log",
	}
	log.SetOutput(&lf)
}

func main() {
	args := os.Args
	if len(args) == 1 {
		config := process.VerifyConfig()
		args = process.TransferConfig(config)
	}

	args = args[1:]
	if len(args)%2 != 0 {
		fmt.Printf("命令行参数解析失败")
		os.Exit(0)
	}

	runtime.GOMAXPROCS(runtime.NumCPU())
	fmt.Printf("当前主机有cpu核数：%v\n", runtime.NumCPU())

	proc := process.ParseParams(args)
	start := time.Now()
	var cost time.Duration = 0.0
	interval := 1

	for {
		innerStart := time.Now()
		wg := sync.WaitGroup{}

		wg.Add(proc.Connections * proc.VisitCounts)
		//传递waitgroup时，需要使用指针传递
		proc.Call(&wg)

		wg.Wait()

		proc.Resp.ReqTotal = int64(proc.Connections * proc.VisitCounts)
		proc.Resp.Tps = float64(proc.Resp.ReqSuccess) / proc.Resp.Avg
		cost = time.Now().Sub(innerStart)

		fmt.Printf("第[%d]轮请求，请求耗时：%v\n", interval, cost)

		allCost := time.Now().Sub(start).Seconds()

		if allCost >= float64(proc.Times) {
			fmt.Printf("整个请求耗时：%v\n", allCost)
			break
		}

		interval++
	}

	proc.Resp.ReqTotal *= int64(interval)

	process.Output(proc)
}

 tool.go

 工具文件主要完成以下功能

1、输出压测结果

2、验证配置文件是否正确且可用

3、打印帮助信息

package process

import (
	"bytes"
	"fmt"
	"os"
)

// Output 输出压测结果到控制台
func Output(proc *Process) {
	var buff bytes.Buffer
	buff.WriteString(fmt.Sprintf("总计请求次数：%d\n", proc.Resp.ReqTotal))
	buff.WriteString(fmt.Sprintf("请求成功次数：%d\n", proc.Resp.ReqSuccess))
	buff.WriteString(fmt.Sprintf("最大请求时间：%fs\n", proc.Resp.Max))
	buff.WriteString(fmt.Sprintf("最小请求时间：%fs\n", proc.Resp.Min))
	buff.WriteString(fmt.Sprintf("平均请求耗时：%fs\n", proc.Resp.Avg))
	buff.WriteString(fmt.Sprintf("Tps：%f\n", proc.Resp.Tps))
	fmt.Printf(buff.String())
}

// VerifyConfig 验证配置文件
func VerifyConfig() *Config {
	config := LoadConfig("config.ini")
	pc := config.ConfigProcess
	currPath, _ := os.Getwd()
	path := string(currPath) + "/config/config.ini"
	if len(pc.Url) == 0 {
		fmt.Printf("\nError: 请在%s配置文件中配置[Url]，或者按照下边的帮助在命令行输入参数\n\n", path)
		Help()
		return nil
	}
	if len(pc.X) == 0 {
		fmt.Printf("\n\n请在%s配置文件中配置[x]，或者按照下边的帮助在命令行输入参数\n\n", path)
		Help()
		return nil
	}
	if pc.C == 0 {
		pc.C = 1
	}

	if pc.N == 0 {
		pc.C = 1
	}
	return config
}

// Help 输出帮助信息
func Help() {
	var buffer bytes.Buffer
	buffer.WriteString("main.go Usage help\n")
	buffer.WriteString("example: main.go -c 2 -n 1 -url http://domain/user/login/login -d '{\"mobile\":\"6712331123\", \"code\" : 888888, \"appId\":1}' -X post -h \"Content-Type:application/json;key:value\"\n")
	buffer.WriteString("param options: \n")
	buffer.WriteString("-c   开启请求协程(线程)数\n")
	buffer.WriteString("-n   每个协程累计请求次数,每次请求都是异步\n")
	buffer.WriteString("-url 请求url地址\n")
	buffer.WriteString("-x 	请求方式，post or get\n")
	buffer.WriteString("-d   请求参数，content-type:application/json 时，参数需要根json字符串，否则content 使用 a=1&b=2 格式连接\n")
	buffer.WriteString("-t   请求累计执行时间\n")
	fmt.Printf("%v\n", buffer.String())

}

logger.go

1、实现io.Writer([]byte) 接口并将日志输出到文件中

2、增加日志按天拆分

package process

import (
	"fmt"
	"os"
	time2 "time"
)

type LogFile struct {
	Filename string
}

// 定义日志输出目录
func (l *LogFile) Write(p []byte) (int, error) {
	rootPwd, _ := os.Getwd()

	rootPwd = rootPwd + "/log/"

	if err := os.Mkdir("log", 0777); err == nil {
		fmt.Printf("日志文件夹创建成功：%v\n", rootPwd)
	}

	time := time2.Now()

	logPrefix := fmt.Sprintf("%4d-%02d-%02d", time.Year(), time.Month(), time.Day())

	f, err := os.OpenFile(rootPwd+logPrefix+l.Filename, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)

	defer f.Close()

	if err != nil {
		return -1, err
	}

	return f.Write(p)
}

config.go

1、结构体ConfigProcess 保存命令行参数信息

2、使用反射获取到结构体ConfigProcess 中的属性，并将解析出的配置信息写入结构体中

3、将读取出的配置信息 模拟写入到 os.Args 数组中，统一传递给 parse.go

package process

import (
	"fmt"
	"log"
	"reflect"
	"strconv"
	"strings"
)

// 压测配置参数
type ConfigProcess struct {
	Url       string `ini:"url"`       //请求url
	X         string `ini:"x"`         //请求方式 post or get
	C         int64  `ini:"c"`         //开启协程数量
	N         int64  `ini:"n"`         //每个协程请求次数
	H         string `ini:"h"`         //header设置
	D         string `ini:"d"`         //请求参数设置
	T         int64  `ini:"t"`         //持续压测时间
	Assertion string `ini:"assertion"` //结果断言，字符串查找
}

// config配置参数
type Config struct {
	ConfigProcess `ini:"process"`
}

// 默认配置文件
var defaultConfig = "[process]\n;;请求地址\nurl=http://domain/user/login/login\n;;请求方式\nx=post\n;;开启协程数量\nc=3\n;;每个协程请求次数\nn=1\n;;header头\nh='Content-Type:application/json'\n;;请求参数\nd={\"mobile\":\"6712331123\", \"code\" : 888888, \"appId\":1}\n;;累计请求时间\n;t=10\n;;断言成功的配置\nassertion=\"code\":0\n"

func init() {
	lf := &LogFile{
		"log.log",
	}
	log.SetOutput(lf)
}

// LoadConfig 加载配置文件
func LoadConfig(fileName string) *Config {
	var config Config
	parseConfig(fileName, &config)
	return &config
}

// TransferConfig 转换配置文件，将配置文件从结构体转为 os.Args 数组
func TransferConfig(config *Config) []string {
	k := reflect.TypeOf(config.ConfigProcess)
	v := reflect.ValueOf(config.ConfigProcess)
	var args []string
	args = append(args, "main.go")
	for i := 0; i < k.NumField(); i++ {
		//fmt.Printf("%v\n", v.Type().)
		field := k.Field(i)
		value := v.Field(i)
		//fmt.Printf("%v\n", field.Name)
		args = append(args, "-"+strings.ToLower(field.Name))
		switch field.Type.Kind() {
		case reflect.String:
			args = append(args, value.String())
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			args = append(args, fmt.Sprintf("%d", value.Int()))
		}
		//	fmt.Printf("%v\n", args)

	}
	return args
}

// 使用反射解析配置文件到结构体
func parseConfig(fileName string, config interface{}) (err error) {
	t := reflect.TypeOf(config)
	if t.Kind() != reflect.Ptr {
		err = fmt.Errorf("no Ptr")
	}
	if t.Elem().Kind() != reflect.Struct {
		err = fmt.Errorf("no Struct")
	}
	fileConent := ReadConfigFile(fileName)
	splitArr := strings.Split(fileConent, "\n")
	//fmt.Printf("配置：%v\n", splitArr)
	var structName string
	for index, line := range splitArr {
		if len(line) == 0 {
			continue
		}

		line = strings.TrimSpace(line)

		if strings.HasPrefix(line, ";") || strings.HasPrefix(line, "#") {
			continue
		}

		if strings.HasPrefix(line, "[") {
			if line[0] != '[' || line[len(line)-1] != ']' {
				fmt.Errorf("%d line error\n", index+1)
				log.Fatalf("%d line error\n", index+1)
				return nil
			}

			sectionName := strings.TrimSpace(line[1 : len(line)-1])
			if len(sectionName) == 0 {
				fmt.Errorf("%d line error\n", index+1)
				log.Fatalf("%d line error\n", index+1)
				return nil
			}
			for i := 0; i < t.Elem().NumField(); i++ {
				field := t.Elem().Field(i)
				if sectionName == field.Tag.Get("ini") {
					structName = field.Name
					break
				}
			}
		} else {
			if strings.Index(line, "=") == -1 || strings.HasPrefix(line, "=") {
				fmt.Errorf("DATA:%d error", index+1)
				log.Fatalf("DATA:%d error", index+1)
				continue
			}
			index := strings.Index(line, "=")
			key := strings.TrimSpace(line[:index])
			value := strings.TrimSpace(line[index+1:])
			v := reflect.ValueOf(config)

			sValue := v.Elem().FieldByName(structName)
			sType := sValue.Type()

			if sType.Kind() != reflect.Struct {
				fmt.Errorf("%s is not struct\n", structName)
				log.Fatalf("%s is not struct\n", structName)
				continue
			}

			var fileNames string
			var fileType reflect.StructField

			for i := 0; i < sValue.NumField(); i++ {
				field := sType.Field(i)
				fileType = field
				if field.Tag.Get("ini") == key {
					fileNames = field.Name
					break
				}
			}

			if len(fileNames) == 0 {
				continue
			}

			fileObj := sValue.FieldByName(fileNames)

			switch fileType.Type.Kind() {
			case reflect.String:
				fileObj.SetString(value)
			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
				parseInt, err := strconv.ParseInt(value, 10, 64)
				if err != nil {
					fmt.Errorf("%s-%s-%s transfer error: %v\n", structName, key, value, err)
					log.Fatalf("%s-%s-%s transfer error: %v\n", structName, key, value, err)
					continue
				}
				fileObj.SetInt(parseInt)
			}
		}

	}
	return err
}

futil.go

1、从config.ini 文件解析配置文件，如果config文件夹不存在则创建,config.ini文件不存在，也会创建，并将默认配置写入 config.ini文件

package process

import (
	"fmt"
	"io"
	"log"
	"os"
)

func init() {
	lf := &LogFile{
		"log.log",
	}
	log.SetOutput(lf)
}

// ReadConfigFile 读取配置文件
func ReadConfigFile(fileName string) string {
	projectPath, _ := os.Getwd()
	projectPath = string(projectPath) + "/config/"

	if err := os.Mkdir(projectPath, 0777); err == nil {
		fmt.Printf("配置文件创建成功\n")
	}

	filePaths := projectPath + fileName

	file, err := os.Open(filePaths)
	defer file.Close()
	if err != nil {
		fmt.Printf("配置文件打开失败: %v, %v\n", projectPath, err)
		log.Printf("配置文件打开失败: %v, %v\n", projectPath, err)
		//return ""
	}

	var chunk []byte

	buf := make([]byte, 1024)

	for {
		n, err := file.Read(buf)
		if err != nil && err != io.EOF {
			fmt.Printf("read buf fail: %v\n", err)
			log.Printf("read buf fail: %v\n", err)

			files, err := os.OpenFile(filePaths, os.O_CREATE|os.O_WRONLY|os.O_RDONLY, 0666)
			defer files.Close()
			if err != nil {
				fmt.Printf("配置文件[%s]不存在,打开失败: %v\n", filePaths, err)
				log.Printf("配置文件[%s]不存在,打开失败: %v\n", filePaths, err)
			}
			files.WriteString(defaultConfig)
			return ""
		}

		if n == 0 {
			//fmt.Printf("内容为空\n")
			break
		}

		chunk = append(chunk, buf[:n]...)
	}

	return string(chunk)
}

parse.go

1、定义命令行参数常量

2、解析命令行参数到结构体Process，用于全局访问，涉及到header对的解析

package process

import (
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
)

// 匹配命令行参数
const (
	N   = "-n"   //每个协程请求次数
	C   = "-c"   //开启多少个协程
	URL = "-url" //请求url
	D   = "-d"   //请求参数

	X = "-x" //请求方式 get or post

	H     = "-h" //请求header
	T     = "-t"
	ASERT = "-assertion" //断言
)

const (
	ParamStep = 2
)

// CommandLine 匹配命令行命令
type CommandLine struct {
	CommandArray map[int]*CommandParams
}

// CommandParams 匹配具体参数
type CommandParams struct {
	Option string
	Value  string
}

// HeaderPairs header头部请求
type HeaderPairs struct {
	Key   string
	Value string
}

// Process 保存并发请求参数
type Process struct {
	Connections int                  //并发请求协程数量
	VisitCounts int                  //每个协程请求次数
	Url         string               //请求地址
	Params      string               //请求参数
	HttpMethod  string               //请求方式 post or get
	Header      map[int]*HeaderPairs //请求header
	Times       int                  //请求时间
	procChan    chan int
	Resp        Response //返回
}

func init() {
	lf = LogFile{
		"log.log",
	}
	log.SetOutput(&lf)
}

// ParseParams 拆分命令行参数
func ParseParams(args []string) (procRes *Process) {

	command := make(map[int]*CommandParams)
	var i = 0
	for p := 0; p < len(args); p = p + ParamStep {
		//每两个参数拆分为一对
		child := args[p : p+ParamStep]
		params := new(CommandParams)
		params.Option = child[0]
		params.Value = child[1]
		command[i] = params
		i++
	}

	proc := new(Process)
	for _, cp := range command {
		if cp == nil {
			continue
		}
		op := strings.ToLower(cp.Option)
		switch op {
		case N:
			proc.VisitCounts, _ = strconv.Atoi(cp.Value)
		case C:
			proc.Connections, _ = strconv.Atoi(cp.Value)
		case URL:
			proc.Url = cp.Value
		case D:
			proc.Params = cp.Value
		case X:
			proc.HttpMethod = strings.ToUpper(cp.Value)
		case H: //解析header
			proc.Header = headerParse(cp.Value)
		case T: //执行时间
			proc.Times, _ = strconv.Atoi(cp.Value)
		case ASERT: //断言结果
			proc.Resp.Assertion = cp.Value
		default: //处理默认值
		}
	}
	if len(proc.Url) == 0 {
		log.Printf("Url 为空，参数解析失败")
		os.Exit(0)
	}
	if len(proc.HttpMethod) == 0 {
		proc.HttpMethod = http.MethodGet
	}
	if proc.VisitCounts == 0 {
		proc.VisitCounts = 1
	}
	if proc.Connections == 0 {
		proc.Connections = 1
	}

	return proc
}

// 解析header
func headerParse(value string) map[int]*HeaderPairs {
	headers := make(map[int]*HeaderPairs)
	value = strings.Trim(value, " ")
	value = strings.ReplaceAll(value, "；", "")
	value = strings.ReplaceAll(value, "：", "")
	var headerValue [100]string
	//var currHeaderKey int
	for k, v := range strings.Split(value, ";") {
		pairs := new(HeaderPairs)
		if len(v) == 0 {
			continue
		}
		d := strings.Split(v, ":")

		pairs.Key = d[0]
		pairs.Value = d[1]
		headerValue[k] = d[1]
		//currHeaderKey = k
		headers[k] = pairs
	}
	//下面这段代码可以忽略,仅仅想把所有请求都加上 application/json
	//var flagType bool
	//for _, vs := range headerValue {
	//	if vs == ContentTypeJson {
	//		flagType = true
	//	}
	//}
	//if flagType == false {
	//	pairs := new(HeaderPairs)
	//	pairs.Key = "Content-Type"
	//	pairs.Value = ContentTypeJson
	//	headers[currHeaderKey+1] = pairs
	//}

	return headers
}

request.go

1、调用请求，传递sync.WaitGroup，并创建2个 channel，用于并发访问和原子接收返回结果

2、传递channel 使用时需要注意两点

2.1 需要将channel的引用传递给每个方法

2.2 通道的发送和接收都需要是协程异步的，否则会报锁错误

3、每个请求执行完成后需要调用 sysn.WaitGroup的Done方法来减计数

package process

import (
	"log"
	"net/http"
	"sync"
)

func init() {
	lf = LogFile{
		"log.log",
	}
	log.SetOutput(&lf)
}

// Call 调用入口方法
func (p *Process) Call(controls *sync.WaitGroup) {
	procChan := make(chan int, p.Connections) //创建有容量的 通道
	respChan := make(chan *Response)
	go p.startCoroutines(procChan)
	go p.startRequest(procChan, controls, respChan)
}

// 创建通道
func (p *Process) startCoroutines(procChan chan int) {
	log.Printf("创建协程数量：%d\n", p.Connections)
	for i := 1; i <= p.Connections; i++ {
		procChan <- i
	}
}

// 消费通道
func (p *Process) startRequest(procChan chan int, controls *sync.WaitGroup, respChan chan *Response) {
	for {
		//println("接收")
		_ = <-procChan
		//fmt.Printf("协程开启第：%d\n", cor)
		for i := 1; i <= p.VisitCounts; i++ {
			go p.request(i, controls, respChan)
		}
	}
}

// 开启协程并发请求
func (p *Process) request(cor int, controls *sync.WaitGroup, respChan chan *Response) {

	defer controls.Done()
	defer func() {
		if err := recover(); err != nil {
			log.Printf("Work failed with %v \n", err)
		}
	}()
	if p.HttpMethod == http.MethodGet {
		p.HttpGet(respChan)
	} else if p.HttpMethod == http.MethodPost {
		p.HttpPostJson(respChan)
	}

}

httprequest.go

1、发送post or get 请求，返回结果时需要记录 最大响应时间，最小响应时间，然后将请求交给repsonse.go去原子统计

package process

import (
	"bytes"
	"io"
	"log"
	"net/http"
	time2 "time"
)

const (
	ContentTypeJson = "application/json"
)

var lf LogFile

// 保存首次请求的耗时，用户最小值判断
var firstCost float64 = 0

func init() {
	lf = LogFile{
		"log.log",
	}
	log.SetOutput(&lf)
}

// HttpPostJson http post请求
func (p *Process) HttpPostJson(respChan chan *Response) {
	var reqParams []byte
	if len(p.Params) >= 0 {
		reqParams = []byte(p.Params)
	}

	reqBody := bytes.NewReader(reqParams)

	request, err := http.NewRequest(http.MethodPost, p.Url, reqBody)
	//增加header
	for _, hv := range p.Header {
		request.Header.Set(hv.Key, hv.Value)
	}
	start := time2.Now()
	client := &http.Client{}
	response, err := client.Do(request)
	if err != nil {
		log.Printf("请求出现错误：%v\n", err)
		return
	}
	resp, _ := io.ReadAll(response.Body)
	cost := time2.Now().Sub(start).Seconds()
	p.Resp.Cost = cost
	p.Resp.Result = string(resp)

	if cost > p.Resp.Max {
		p.Resp.Max = cost
	}

	if firstCost == 0 {
		firstCost = cost
		p.Resp.Min = cost
	}

	if cost < p.Resp.Min {
		p.Resp.Min = cost
	}

	//使用通道串行处理返回结果
	go p.Resp.Receive(respChan)
	go p.Resp.getResult(respChan)

	log.Printf("返回结果：%v\n", string(resp))
}

// HttpGet http get请求
func (p *Process) HttpGet(respChan chan *Response) {
	start := time2.Now()
	response, err := http.Get(p.Url)
	if err != nil {
		log.Printf("请求链接有误: %v\n", err)
		return
	}

	resp, _ := io.ReadAll(response.Body)

	cost := time2.Now().Sub(start).Seconds()
	p.Resp.Cost = cost
	p.Resp.Result = string(resp)
	if cost > p.Resp.Max {
		p.Resp.Max = cost
	}

	if firstCost == 0 {
		firstCost = cost
		p.Resp.Min = cost
	}

	if cost < p.Resp.Min {
		p.Resp.Min = cost
	}

	//使用通道串行处理返回结果
	go p.Resp.Receive(respChan)
	go p.Resp.getResult(respChan)

	log.Printf("get返回结果：%v\n", string(resp))
}

response.go

1、定义结构体Response 用于保存 压测结果 和响应返回信息

2、通过channel 为1 来原子执行go的协程返回

3、使用 rest 全局变量来保存 请求响应时间，用来计算平均响应时间

4、断言默认值是 "\"code\":0"， 可以在配置文件assertion 参数来设置

package process

import (
	"strings"
)

// Response 保存返回的数据，压测完后保存结果
type Response struct {
	Result     string  //请求返回结果
	Cost       float64 //每次请求耗时
	Max        float64 //最大请求时间
	Min        float64 //最小请求时间
	Avg        float64 //平均请求时间
	Tps        float64 //吞吐量
	ReqTotal   int64   //请求次数
	ReqSuccess int64   //成功次数
	Assertion  string  //存放断言

}

// 切片保存每次请求的耗时
var rest []float64

// Receive 通道接收返回
func (r *Response) Receive(cost chan *Response) {
	cost <- r
}

// 通道处理返回
func (p *Response) getResult(cost chan *Response) {
	for {
		ct := <-cost
		rest = append(rest, ct.Cost)
		sum := 0.0
		total := 0.0
		for _, num := range rest {
			sum += num
			total += 1
		}
		//默认断言
		var asserts = "\"code\":0"
		if len(p.Assertion) > 0 {
			asserts = p.Assertion
		}
		if strings.Contains(ct.Result, asserts) {
			p.ReqSuccess++
		}
		p.Avg = sum / total
	}
}