基于Go语言实现一个压测工具

// request 请求数据
type request struct {
	url       string            // url
	form      string            // http/websocket/tcp
	method    string            // 方法 get/post/put
	headers   map[string]string // headers
	body      string            // body
	verify    string            // 验证的方法
	timeout   time.duration     // 请求超时时间
	debug     bool              // 是否开启debug模式
	maxcon    int               // 每个连接的请求数
	http2     bool              // 是否使用http2.0
	keepalive bool              // 是否开启长连接
	code      int               // 验证的状态码
	redirect  bool              // 是否重定向
}

// newrequest 生成请求结构体
// url 压测的url
// verify 验证方法 在server/verify中 http 支持:statuscode、json websocket支持:json
// timeout 请求超时时间
// debug 是否开启debug
// path curl文件路径 http接口压测，自定义参数设置
func newrequest(url string, verify string, code int, timeout time.duration, debug bool, path string,
	reqheaders []string, reqbody string, maxcon int, http2, keepalive, redirect bool) (request *request, err error) {
	var (
		method  = "get"
		headers = make(map[string]string)
		body    string
	)
	if path != "" {
		var curl *curl
		curl, err = parsethefile(path)
		if err != nil {
			return nil, err
		}
		if url == "" {
			url = curl.geturl()
		}
		method = curl.getmethod()
		headers = curl.getheaders()
		body = curl.getbody()
	} else {
		if reqbody != "" {
			method = "post"
			body = reqbody
		}
		for _, v := range reqheaders {
			getheadervalue(v, headers)
		}
		if _, ok := headers["content-type"]; !ok {
			headers["content-type"] = "application/x-www-form-urlencoded; charset=utf-8"
		}
	}
	var form string
	form, url = getform(url)
	if form == "" {
		err = fmt.errorf("url:%s 不合法,必须是完整http、websocket连接", url)
		return
	}
	var ok bool
	switch form {
	case formtypehttp:
		// verify
		if verify == "" {
			verify = "statuscode"
		}
		key := fmt.sprintf("%s.%s", form, verify)
		_, ok = verifymaphttp[key]
		if !ok {
			err = errors.new("验证器不存在:" + key)
			return
		}
	case formtypewebsocket:
		// verify
		if verify == "" {
			verify = "json"
		}
		key := fmt.sprintf("%s.%s", form, verify)
		_, ok = verifymapwebsocket[key]
		if !ok {
			err = errors.new("验证器不存在:" + key)
			return
		}
	}
	if timeout == 0 {
		timeout = 30 * time.second
	}
	request = &request{
		url:       url,
		form:      form,
		method:    strings.toupper(method),
		headers:   headers,
		body:      body,
		verify:    verify,
		timeout:   timeout,
		debug:     debug,
		maxcon:    maxcon,
		http2:     http2,
		keepalive: keepalive,
		code:      code,
		redirect:  redirect,
	}
	return
}

// requestresults 请求结果
type requestresults struct {
	id            string // 消息id
	chanid        uint64 // 消息id
	time          uint64 // 请求时间 纳秒
	issucceed     bool   // 是否请求成功
	errcode       int    // 错误码
	receivedbytes int64
}

var (
    mutex sync.rwmutex
    // clients 客户端
    // key 客户端id - value 客户端
    clients = make(map[uint64]*http.client)
)

// createlanghttpclient 初始化长连接客户端参数
// 创建了一个配置了长连接的 http 客户端传输对象
func createlanghttpclient(request *model.request) *http.client {
    tr := &http.transport{
        // 使用 net.dialer 来建立 tcp 连接
        // timeout 设置为 30 秒，表示如果连接在 30 秒内没有建立成功，则超时
        // keepalive 设置为 30 秒，表示连接建立后，如果 30 秒内没有数据传输，则发送一个 keep-alive 探测包以保持连接
        dialcontext: (&net.dialer{
            timeout:   30 * time.second,
            keepalive: 30 * time.second,
        }).dialcontext,
        maxidleconns:        0,                // 最大连接数,默认0无穷大
        maxidleconnsperhost: request.maxcon,   // 对每个host的最大连接数量(maxidleconnsperhost<=maxidleconns)
        idleconntimeout:     90 * time.second, // 多长时间未使用自动关闭连接
        // insecureskipverify 设置为 true，表示不验证服务器的 ssl 证书
        tlsclientconfig: &tls.config{insecureskipverify: true},
    }
    if request.http2 {
        // 使用真实证书 验证证书 模拟真实请求
        tr = &http.transport{
            dialcontext: (&net.dialer{
                timeout:   30 * time.second,
                keepalive: 30 * time.second,
            }).dialcontext,
            maxidleconns:        0,                // 最大连接数,默认0无穷大
            maxidleconnsperhost: request.maxcon,   // 对每个host的最大连接数量(maxidleconnsperhost<=maxidleconns)
            idleconntimeout:     90 * time.second, // 多长时间未使用自动关闭连接
            // 配置 tls 客户端设置，insecureskipverify 设置为 false，表示验证服务器的 ssl 证书
            tlsclientconfig: &tls.config{insecureskipverify: false},
        }
        // 将 tr 配置为支持 http/2 协议
        _ = http2.configuretransport(tr)
    }

    client := &http.client{
        transport: tr,
    }

    // 禁止 http 客户端自动重定向，而是让客户端在遇到重定向时停止并返回最后一个响应
    if !request.redirect {
        client.checkredirect = func(req *http.request, via []*http.request) error {
            return http.erruselastresponse
        }
    }

​​​​​​​    return client
}

// grpcsocket grpc
type grpcsocket struct {
    conn    *grpc.clientconn
    address string
}

// websocket websocket
type websocket struct {
    conn       *websocket.conn
    urllink    string
    url        *url.url
    isssl      bool
    httpheader map[string]string
}

// grpcrequest 请求
func grpcrequest(chanid uint64, ch chan<- *model.requestresults, i uint64, request *model.request,
    ws *client.grpcsocket) {
    var (
        starttime = time.now()
        issucceed = false
        errcode   = model.httpok
    )
    // 获取连接
    conn := ws.getconn()
    if conn == nil {
        errcode = model.requesterr
    } else {
        c := pb.newapiserverclient(conn)
        var (
            ctx = context.background()
            req = &pb.request{
                username: request.body,
            }
        )
        // 发送请求，获得响应
        rsp, err := c.helloworld(ctx, req)
        if err != nil {
            errcode = model.requesterr
        } else {
            // 200 为成功
            if rsp.code != 200 {
                errcode = model.requesterr
            } else {
                issucceed = true
            }
        }
    }
    requesttime := uint64(helper.diffnano(starttime))
    requestresults := &model.requestresults{
        time:      requesttime,
        issucceed: issucceed,
        errcode:   errcode,
    }
    requestresults.setid(chanid, i)
    ch <- requestresults
}

// http 请求
func http(ctx context.context, chanid uint64, ch chan<- *model.requestresults, totalnumber uint64, wg *sync.waitgroup,
	request *model.request) {
	defer func() {
		wg.done()
	}()
	for i := uint64(0); i < totalnumber; i++ {
		if ctx.err() != nil {
			break
		}

		list := getrequestlist(request)
		issucceed, errcode, requesttime, contentlength := sendlist(chanid, list)
		requestresults := &model.requestresults{
			time:          requesttime,
			issucceed:     issucceed,
			errcode:       errcode,
			receivedbytes: contentlength,
		}
		requestresults.setid(chanid, i)
		ch <- requestresults
	}

	return
}

// calculatedata 计算数据
func calculatedata(concurrent, processingtime, requesttime, maxtime, mintime, successnum, failurenum uint64,
	chanidlen int, errcode *sync.map, receivedbytes int64) {
	if processingtime == 0 {
		processingtime = 1
	}
	var (
		qps              float64
		averagetime      float64
		maxtimefloat     float64
		mintimefloat     float64
		requesttimefloat float64
	)
	// 平均 qps 成功数*总协程数/总耗时 (每秒)
	if processingtime != 0 {
		qps = float64(successnum*concurrent) * (1e9 / float64(processingtime))
	}
	// 平均时长 总耗时/总请求数/并发数 纳秒=>毫秒
	if successnum != 0 && concurrent != 0 {
		averagetime = float64(processingtime) / float64(successnum*1e6)
	}
	// 纳秒=>毫秒
	maxtimefloat = float64(maxtime) / 1e6
	mintimefloat = float64(mintime) / 1e6
	requesttimefloat = float64(requesttime) / 1e9
	// 打印的时长都为毫秒
	table(successnum, failurenum, errcode, qps, averagetime, maxtimefloat, mintimefloat, requesttimefloat, chanidlen,
		receivedbytes)
}