设计一个简单的压力测试工具

package worker

import "fmt"

type worker struct {
	url           string
	totalReqNum   int
	concurrentNum int
	jobsCh        chan struct{}
	resultCh      chan int64
}

func NewWorker(url string, concurrentNum int, totalReqNum int) *worker {
	return &worker{
		url:           url,
		concurrentNum: concurrentNum,
		totalReqNum:   totalReqNum,
		jobsCh:        make(chan struct{}, totalReqNum),
		resultCh:      make(chan int64, totalReqNum),
	}
}

type WorkFunc interface {
	DoWork() int64
}

func (w *worker) BuildWorker(wf WorkFunc) {
	for i := 1; i <= w.concurrentNum; i++ {
		go doWork(wf, w.jobsCh, w.resultCh)
		//fmt.Println("worker ", i, " initialized")
	}
}

func (w *worker) BuildJobs() {
	for i := 0; i < w.totalReqNum; i++ {
		w.jobsCh <- struct{}{}
		//fmt.Println("add job ", i+1)
	}
}

func (w *worker) PrintStatistic() {
	totalRespTime := int64(0)
	nfpRespTime := int64(0)
	nfpCount := int(float64(w.totalReqNum) * 0.95)
	for i := 0; i < w.totalReqNum; i++ {
		t := <-w.resultCh
		totalRespTime += t
		if nfpCount >= i {
			nfpRespTime += t
		}
	}

	fmt.Println("")
	fmt.Printf("avg response time:\t%.2Fs\n", calcRespTime(totalRespTime, w.totalReqNum))
	fmt.Printf("95%% response time:\t%.2Fs\n", calcRespTime(nfpRespTime, nfpCount))
	fmt.Println("")
}

func calcRespTime(totalNanoTime int64, totalCount int) float64 {
	return float64(totalNanoTime)/float64(totalCount)/float64(1000000000)
}

func doWork(wf WorkFunc, jobs <-chan struct{}, respTimeCh chan<- int64) {
	for range jobs {
		respTime := wf.DoWork()
		respTimeCh <- respTime
		//fmt.Println("resp time: ", respTime)
	}
}

package worker

import (
	"github.com/magiconair/properties/assert"
	"testing"
	"time"
)

type mockWorkFunc struct {
	now int64
}

func (m *mockWorkFunc) DoWork() int64 {
	return m.now
}

func TestWorker_BuildWorker(t *testing.T) {
	total := 1
	w := &worker{
		url:           "http://www.example.com",
		totalReqNum:   total,
		concurrentNum: 1,
		jobsCh:        make(chan struct{}, total),
		resultCh:      make(chan int64, total),
	}
	now := time.Now().Unix()

	w.BuildWorker(&mockWorkFunc{now: now})

	w.jobsCh <- struct{}{}

	res := <-w.resultCh
	assert.Equal(t, res, now)
}

func TestWorker_BuildJobs(t *testing.T) {
	total := 5
	w := &worker{
		url:           "http://www.example.com",
		totalReqNum:   total,
		concurrentNum: 1,
		jobsCh:        make(chan struct{}, total),
		resultCh:      make(chan int64, total),
	}

	w.BuildJobs()

	var count int

	for range w.jobsCh {
		count++
		if count == total {
			break
		}
	}

	assert.Equal(t, count, total)
}

func TestWorker_PrintStatistic(t *testing.T) {
	total := 5
	w := &worker{
		url:           "http://www.example.com",
		totalReqNum:   total,
		concurrentNum: 1,
		jobsCh:        make(chan struct{}, total),
		resultCh:      make(chan int64, total),
	}

	now := time.Now().Unix()
	w.BuildWorker(&mockWorkFunc{now: now})
	w.BuildJobs()
	w.PrintStatistic()
}

package worker

import (
	"fmt"
	"net/http"
	"time"
)

type urlWorkFunc struct {
	url string
}

func NewUrlWorkFunc(url string) WorkFunc {
	return &urlWorkFunc{url: url}
}

func (u *urlWorkFunc) DoWork() int64 {
	start := time.Now().UnixNano()
	_, err := http.Get(u.url)
	if err != nil {
		fmt.Println(err)
	}
	//time.Sleep(1 * time.Second)
	end := time.Now().UnixNano()

	return end - start
}

package worker

import (
	"github.com/magiconair/properties/assert"
	"testing"
)

func Test_DoWork(t *testing.T) {
	u := &urlWorkFunc{url: "https://www.baidu.com"}
	latency := u.DoWork()
	assert.Equal(t, latency >= 0, true)
}

package main

import (
	"flag"
	"fmt"
	"net/url"
	"pressure-test-toy/pkg/worker"
)

var (
	targetURL     = flag.String("url", "", "target URL for pressure test")
	concurrentNum = flag.Int("concurrentNum", 1, "concurrency number")
	totalReqNum   = flag.Int("totalReqNum", 1, "total request number")
)

func main() {
	flag.Parse()
	if *targetURL == "" {
		flag.Usage()
		return
	}

	_, err := url.ParseRequestURI(*targetURL)
	if err != nil {
		fmt.Println("invalid target url: ", *targetURL)
		return
	}

	u := worker.NewUrlWorkFunc(*targetURL)
	w := worker.NewWorker(*targetURL, *concurrentNum, *totalReqNum)

	w.BuildWorker(u)
	w.BuildJobs()
	w.PrintStatistic()
}

#!make

# version
version := "v0.0.1"

# app name
app := "pttoy"

# platform
darwin := ${app}-darwin.${version}.bin
linux := ${app}-linux.${version}.bin

default: darwin

linux:

	CGO_ENABLE=0 GOOS=linux GOARCH=amd64 go build -o "${linux}" cmd/main.go
	upx ${linux}

darwin:

	CGO_ENABLE=0 GOOS=darwin GOARCH=amd64 go build -o "${darwin}" cmd/main.go
	upx ${darwin}