fix: unittest

docs/cmd/hrp_boom.md

@@ -27,9 +27,10 @@ hrp boom [flags]
       --disable-compression             Disable compression
       --disable-console-output          Disable console output.
       --disable-keepalive               Disable keepalive
-      --expect-workers int              How many workers master should expect to connect before starting the test (only when --autostart is used (default 1)
+      --expect-workers int              How many workers master should expect to connect before starting the test (only when --autostart is used) (default 1)
       --expect-workers-max-wait int     How many workers master should expect to connect before starting the test (only when --autostart is used
   -h, --help                            help for boom
+      --ignore-quit                     ignores quit from master (only when --worker is used)
       --loop-count int                  The specify running cycles for load testing (default -1)
       --master                          master of distributed testing
       --master-bind-host string         Interfaces (hostname, ip) that hrp master should bind to. Only used when running with --master. Defaults to * (all available interfaces). (default "127.0.0.1")

hrp/internal/boomer/runner.go

@@ -151,9 +151,9 @@ func (c *Controller) increaseFinishedCount() {
 func (c *Controller) reset() {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
-	c.spawnCount = 0
+	atomic.StoreInt64(&c.spawnCount, 0)
 	c.spawnRate = 0
-	c.currentClientsNum = 0
+	atomic.StoreInt64(&c.currentClientsNum, 0)
 	c.spawnDone = make(chan struct{})
 	c.tasks = []*Task{}
 	c.once = sync.Once{}
@@ -490,6 +490,7 @@ func (r *runner) stop() {
 	if r.rateLimitEnabled {
 		r.rateLimiter.Stop()
 	}
+	r.updateState(StateStopped)
 }
 
 func (r *runner) getState() int32 {
@@ -581,6 +582,7 @@ type workerRunner struct {
 
 	tasksChan chan *profileMessage
 
+	mutex      sync.Mutex
 	ignoreQuit bool
 }
 
@@ -597,6 +599,7 @@ func newWorkerRunner(masterHost string, masterPort int) (r *workerRunner) {
 		masterPort: masterPort,
 		nodeID:     getNodeID(),
 		tasksChan:  make(chan *profileMessage, 10),
+		mutex:      sync.Mutex{},
 	}
 	return r
 }
@@ -746,6 +749,8 @@ func (r *workerRunner) run() {
 
 // start load test
 func (r *workerRunner) start() {
+	r.mutex.Lock()
+	defer r.mutex.Unlock()
 	r.reset()
 
 	// start rate limiter
@@ -768,7 +773,6 @@ func (r *workerRunner) stop() {
 	if r.isStarted() {
 		r.runner.stop()
 		close(r.rebalance)
-		r.updateState(StateStopped)
 	}
 }
 

hrp/internal/builtin/utils.go

@@ -1,13 +1,11 @@
 package builtin
 
 import (
-	"archive/zip"
 	"bufio"
 	"bytes"
 	"encoding/csv"
 	builtinJSON "encoding/json"
 	"fmt"
-	"io"
 	"math/rand"
 	"os"
 	"os/exec"
@@ -493,156 +491,6 @@ func GetFileNameWithoutExtension(path string) string {
 	return base[0 : len(base)-len(ext)]
 }
 
-func ZipDir(filename string, root string) error {
-	p, err := os.Getwd()
-	if err != nil {
-		return err
-	}
-	if strings.Contains(root, p) {
-		root, err = filepath.Rel(p, root)
-		if err != nil {
-			return err
-		}
-	}
-	err = os.RemoveAll(filename)
-	if err != nil {
-		return err
-	}
-	var files []string
-	err = filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
-		files = append(files, path)
-		return nil
-	})
-	if err != nil {
-		return err
-	}
-	err = ZipFiles(filename, files)
-	return err
-}
-
-// ZipFiles compresses one or many files into a single zip archive file.
-// Param 1: filename is the output zip file's name.
-// Param 2: files is a list of files to add to the zip.
-func ZipFiles(filename string, files []string) error {
-	newZipFile, err := os.Create(filename)
-	if err != nil {
-		return err
-	}
-	defer newZipFile.Close()
-
-	zipWriter := zip.NewWriter(newZipFile)
-	defer zipWriter.Close()
-
-	// Add files to zip
-	for _, file := range files {
-		if err = AddFileToZip(zipWriter, file); err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-func AddFileToZip(zipWriter *zip.Writer, filename string) error {
-	fileToZip, err := os.Open(filename)
-	if err != nil {
-		return err
-	}
-	defer fileToZip.Close()
-
-	// Get the file information
-	info, err := fileToZip.Stat()
-	if err != nil {
-		return err
-	}
-
-	header, err := zip.FileInfoHeader(info)
-	if err != nil {
-		return err
-	}
-
-	// Using FileInfoHeader() above only uses the basename of the file. If we want
-	// to preserve the folder structure we can overwrite this with the full path.
-	header.Name = filename
-
-	// if dir
-	if info.IsDir() {
-		header.Name += `/`
-	} else {
-		// Change to deflate to gain better compression
-		// see http://golang.org/pkg/archive/zip/#pkg-constants
-		header.Method = zip.Deflate
-	}
-
-	writer, err := zipWriter.CreateHeader(header)
-	if err != nil {
-		return err
-	}
-	if !info.IsDir() {
-		_, err = io.Copy(writer, fileToZip)
-	}
-	return err
-}
-
-func UnZip(dst, src string) (err error) {
-	zr, err := zip.OpenReader(src)
-	defer zr.Close()
-	if err != nil {
-		return
-	}
-	if dst != "" {
-		if err := os.MkdirAll(dst, 0755); err != nil {
-			return err
-		}
-	}
-	for _, file := range zr.File {
-		path := filepath.Join(dst, file.Name)
-		if file.FileInfo().IsDir() {
-			if err := os.MkdirAll(path, file.Mode()); err != nil {
-				return err
-			}
-			continue
-		}
-		fr, err := file.Open()
-		if err != nil {
-			return err
-		}
-		fw, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR|os.O_TRUNC, file.Mode())
-		if err != nil {
-			return err
-		}
-		_, err = io.Copy(fw, fr)
-		if err != nil {
-			return err
-		}
-		log.Info().Msg(fmt.Sprintf("unzip %s successful\n", path))
-		_ = fw.Close()
-		_ = fr.Close()
-	}
-	return nil
-}
-
-func File2Bytes(filename string) ([]byte, error) {
-	file, err := os.Open(filename)
-	if err != nil {
-		return nil, err
-	}
-	defer file.Close()
-
-	stats, err := file.Stat()
-	if err != nil {
-		return nil, err
-	}
-
-	data := make([]byte, stats.Size())
-	count, err := file.Read(data)
-	if err != nil {
-		return nil, err
-	}
-	log.Info().Msg(fmt.Sprintf("read file %s len: %d \n", filename, count))
-
-	return data, nil
-}
-
 func Bytes2File(data []byte, filename string) error {
 	file, err := os.Create(filename)
 	if err != nil {