Merge pull request #7 from ryo-yamaoka/add-max-rps

Add max RPS option

README.md

@@ -21,20 +21,20 @@ See [./sample/main.go](./sample/main.go) for a sample.
 $ go run ./sample/...
 
 [Setting]
-* warm up time:   5s
-* duration:       1s
-* max concurrent: 1
+* warm up time:   2s
+* duration:       3s
+* max concurrent: 2
 
 [Request]
-* total:      90
-* succeeded:  90
+* total:      7
+* succeeded:  7
 * failed:     0
 * error rate: 0 %
-* RPS:        90
+* RPS:        2.3
 
 [Latency]
-* max: 11.0 ms
-* min: 10.0 ms
+* max: 11.1 ms
+* min: 10.6 ms
 * avg: 10.9 ms
 * med: 11.0 ms
 * 99th percentile: 11.0 ms
@@ -47,9 +47,10 @@ When you useing `otchkiss.New()` or `setting.FromDefaultFlag()`, will be parsed
 
 This eliminates the need to write the parsing process.
 
-* `-p`: Specify the number of parallels executions (default: `1`, it's not concurrently)
-* `-d`: Running duration, ex: 300s or 5m etc... (default: `1s`)
+* `-p`: Specify the number of parallels executions. `0` means unlimited (default: `1`, it's not concurrently)
+* `-d`: Running duration, ex: 300s or 5m etc... (default: `5s`)
 * `-w`: Exclude from results for a given time after startup, ex: 300s or 5m etc... (default: `5s`)
+* `-r`: Specify the max request per second. 0 means unlimited (default: `1`)
 
 ## Development
 

go.mod

@@ -6,10 +6,12 @@ require (
 	github.com/dustin/go-humanize v1.0.1
 	github.com/google/go-cmp v0.5.9
 	github.com/stretchr/testify v1.8.0
+	go.uber.org/ratelimit v0.2.0
 	golang.org/x/sync v0.3.0
 )
 
 require (
+	github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect

go.sum

@@ -1,8 +1,8 @@
+github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129 h1:MzBOUgng9orim59UnfUTLRjMpd09C5uEVQ6RPGeCaVI=
+github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129/go.mod h1:rFgpPQZYZ8vdbc+48xibu8ALc3yeyd64IhHS+PU6Yyg=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo=
-github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
 github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
@@ -11,11 +11,15 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
+github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
-golang.org/x/sync v0.0.0-20220819030929-7fc1605a5dde h1:ejfdSekXMDxDLbRrJMwUk6KnSLZ2McaUCVcIKM+N6jc=
-golang.org/x/sync v0.0.0-20220819030929-7fc1605a5dde/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw=
+go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
+go.uber.org/ratelimit v0.2.0 h1:UQE2Bgi7p2B85uP5dC2bbRtig0C+OeNRnNEafLjsLPA=
+go.uber.org/ratelimit v0.2.0/go.mod h1:YYBV4e4naJvhpitQrWJu1vCpgB7CboMe0qhltKt6mUg=
 golang.org/x/sync v0.3.0 h1:ftCYgMx6zT/asHUrPw8BLLscYtGznsLAnjq5RH9P66E=
 golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=

otchkiss.go

@@ -5,14 +5,16 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"sync"
 	"text/template"
 	"time"
 
 	"github.com/ryo-yamaoka/otchkiss/result"
+	"github.com/ryo-yamaoka/otchkiss/sema"
 	"github.com/ryo-yamaoka/otchkiss/setting"
 
 	humanize "github.com/dustin/go-humanize"
-	"golang.org/x/sync/semaphore"
+	"go.uber.org/ratelimit"
 )
 
 // Requester defines the behavior of the request that Otchkiss performs.
@@ -40,9 +42,13 @@ type Otchkiss struct {
 
 // New returns Otchkiss instance with default setting.
 // By default, the following three command line arguments are parsed and set.
-//   -p: Specify the number of parallels executions (default: 1, it's not concurrently)
-//   -d: Running duration, ex: 300s or 5m etc... (default: 1s)
-//   -w: Exclude from results for a given time after startup, ex: 300s or 5m etc... (default: 5s)
+//
+//	-p: Specify the number of parallels executions. 0 means unlimited (default: 1, it's not concurrently)
+//	-d: Running duration, ex: 300s or 5m etc... (default: 5s)
+//	-w: Exclude from results for a given time after startup, ex: 300s or 5m etc... (default: 5s)
+//	-r: Specify the max request per second. 0 means unlimited (default: 1)
+//
+// Note: -p or -r, whichever is smaller blocks the request.
 func New(requester Requester) (*Otchkiss, error) {
 	s, err := setting.FromDefaultFlag()
 	if err != nil {
@@ -82,10 +88,10 @@ func new(requester Requester, setting *setting.Setting, r *result.Result) (*Otch
 }
 
 // Start run Otchkiss load testing, and the test follows these steps.
-//   1. Run Init()
-//   2. Start RequestOne() repeatedly as warm up (it will NOT count as Result)
-//   3. Start RequestOne() repeatedly as actual test (it will count as Result)
-//   4. End RequestOne() execute and run Terminate()
+//  1. Run Init()
+//  2. Start RequestOne() repeatedly as warm up (it will NOT count as Result)
+//  3. Start RequestOne() repeatedly as actual test (it will count as Result)
+//  4. End RequestOne() execute and run Terminate()
 func (ot *Otchkiss) Start(ctx context.Context) error {
 	if err := ot.Requester.Init(); err != nil {
 		return fmt.Errorf("failed to initialize requester: %w", err)
@@ -100,12 +106,26 @@ func (ot *Otchkiss) Start(ctx context.Context) error {
 		close(warmUp)
 	}()
 
-	sem := semaphore.NewWeighted(int64(ot.Setting.MaxConcurrent))
+	sem := sema.NewWeighted(int64(ot.Setting.MaxConcurrent))
+	rl := ratelimit.NewUnlimited()
+	maxRPS := ot.Setting.MaxRPS
+	if maxRPS != 0 {
+		rl = ratelimit.New(maxRPS, ratelimit.Per(1*time.Second))
+	}
+
+	var wg sync.WaitGroup
 	for {
+		if ctx.Err() != nil {
+			break
+		}
 		if err := sem.Acquire(ctx, 1); err != nil {
-			return ot.Requester.Terminate()
+			break
 		}
+		rl.Take()
+
+		wg.Add(1)
 		go func() {
+			defer wg.Done()
 			start := time.Now()
 			err := ot.Requester.RequestOne(ctx)
 			elapsed := time.Since(start) // Do this before error handling to obtain the most accurate time possible.
@@ -124,6 +144,9 @@ func (ot *Otchkiss) Start(ctx context.Context) error {
 			}
 		}()
 	}
+
+	wg.Wait()
+	return ot.Requester.Terminate()
 }
 
 type ReportParams struct {

otchkiss_test.go

@@ -40,7 +40,8 @@ func TestNew(t *testing.T) {
 				Requester: &testRequesterImpl{},
 				Setting: &setting.Setting{
 					MaxConcurrent: 1,
-					RunDuration:   1 * time.Second,
+					MaxRPS:        1,
+					RunDuration:   5 * time.Second,
 					WarmUpTime:    5 * time.Second,
 				},
 			},

sample/main.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/ryo-yamaoka/otchkiss"
+	"github.com/ryo-yamaoka/otchkiss/setting"
 )
 
 type SampleRequester struct{}
@@ -33,7 +34,13 @@ func main() {
 }
 
 func run() error {
-	ot, err := otchkiss.New(&SampleRequester{})
+	st := &setting.Setting{
+		RunDuration:   3 * time.Second,
+		WarmUpTime:    2 * time.Second,
+		MaxConcurrent: 2,
+		MaxRPS:        2,
+	}
+	ot, err := otchkiss.FromConfig(&SampleRequester{}, st, 1_000_000)
 	if err != nil {
 		return fmt.Errorf("init error: %w", err)
 	}

sema/sema.go

@@ -0,0 +1,43 @@
+package sema
+
+import (
+	"context"
+
+	"golang.org/x/sync/semaphore"
+)
+
+// Sema implements semaphore with it can be unlimited by specifying 0
+type Sema struct {
+	sem *semaphore.Weighted
+}
+
+// NewWeighted creates a new weighted semaphore with the given maximum combined weight for concurrent access.
+// When you specify 0, it means unlimited concurrent access.
+func NewWeighted(n int64) *Sema {
+	var sem *semaphore.Weighted
+	if n != 0 {
+		sem = semaphore.NewWeighted(n)
+	}
+	return &Sema{
+		sem: sem,
+	}
+}
+
+// Acquire acquires the semaphore with a weight of n, blocking until resources are available or ctx is done. On success, returns nil. On failure, returns ctx.Err() and leaves the semaphore unchanged.
+// If ctx is already done, Acquire may still succeed without blocking.
+// If you created semaphore with 0, this method is non blocking.
+func (s *Sema) Acquire(ctx context.Context, n int64) error {
+	if s.sem == nil {
+		return ctx.Err()
+	}
+	return s.sem.Acquire(ctx, n)
+}
+
+// Release releases the semaphore with a weight of n.
+// But if you created semaphore with 0, this method is no effect.
+func (s *Sema) Release(n int64) {
+	if s.sem == nil {
+		return
+	}
+	s.sem.Release(n)
+}

setting/setting.go

@@ -9,12 +9,15 @@ import (
 
 const (
 	defaultMaxConcurrent = 1
-	defaultRunDuration   = 1 * time.Second
+	defaultRunDuration   = 5 * time.Second
 	defaultWarmUpTime    = 5 * time.Second
+	defaultMaxRPS        = 1
 )
 
 type Setting struct {
 	// MaxConcurrent defines how many RequestOne should concurrently running.
+	// 0 means unlimited.
+	// MaxConcurrent or MaxRPS, whichever is smaller blocks the request.
 	MaxConcurrent int
 
 	// RunDuration defines how long RequestOne should continue to run.
@@ -24,29 +27,38 @@ type Setting struct {
 	// WarmUpTime defines how long RequestOne not included in the measurement should continue to run.
 	// During the time specified here, the request results are NOT included in the Result.
 	WarmUpTime time.Duration
+
+	// MaxRPS defines how much can request per second.
+	// 0 means unlimited.
+	// MaxConcurrent or MaxRPS, whichever is smaller blocks the request.
+	MaxRPS int
 }
 
 // New returns Setting instance made by user defined config.
-func New(maxConcurrent int, runDuration, warmUpTime time.Duration) (*Setting, error) {
-	return newSetting(maxConcurrent, runDuration, warmUpTime)
+func New(maxConcurrent, maxRPS int, runDuration, warmUpTime time.Duration) (*Setting, error) {
+	return newSetting(maxConcurrent, maxRPS, runDuration, warmUpTime)
 }
 
 // FromDefaultConfig returns Setting by flag or default value config.
 func FromDefaultFlag() (*Setting, error) {
 	c := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
-	maxConcurrent := c.Int("p", defaultMaxConcurrent, "Specify the number of parallels executions (default: 1, it's not concurrently)")
-	runDuration := c.Duration("d", defaultRunDuration, "Running duration, ex: 300s or 5m etc... (default: 1s)")
+	maxConcurrent := c.Int("p", defaultMaxConcurrent, "Specify the number of parallels executions. 0 means unlimited (default: 1, it's not concurrently)")
+	runDuration := c.Duration("d", defaultRunDuration, "Running duration, ex: 300s or 5m etc... (default: 5s)")
 	warmUpTime := c.Duration("w", defaultWarmUpTime, "Exclude from results for a given time after startup, ex: 300s or 5m etc... (default: 5s)")
+	maxRPS := c.Int("r", defaultMaxRPS, "Specify the max request per second. 0 means unlimited (default: 1)")
 	if err := c.Parse(os.Args[1:]); err != nil {
 		return nil, err
 	}
 
-	return newSetting(*maxConcurrent, *runDuration, *warmUpTime)
+	return newSetting(*maxConcurrent, *maxRPS, *runDuration, *warmUpTime)
 }
 
-func newSetting(maxConcurrent int, runDuration, warmUpTime time.Duration) (*Setting, error) {
-	if !(maxConcurrent > 0) {
-		return nil, errors.New("max concurrent must be > 0")
+func newSetting(maxConcurrent, maxRPS int, runDuration, warmUpTime time.Duration) (*Setting, error) {
+	if !(maxConcurrent >= 0) {
+		return nil, errors.New("max concurrent must be >= 0")
+	}
+	if !(maxRPS >= 0) {
+		return nil, errors.New("max RPS must be >= 0")
 	}
 	if !(runDuration > 0*time.Second) {
 		return nil, errors.New("run duration must be > 0 sec")
@@ -59,5 +71,6 @@ func newSetting(maxConcurrent int, runDuration, warmUpTime time.Duration) (*Sett
 		MaxConcurrent: maxConcurrent,
 		RunDuration:   runDuration,
 		WarmUpTime:    warmUpTime,
+		MaxRPS:        maxRPS,
 	}, nil
 }