func TestContextCancelWhileControlBatchLimiter(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
limitCtx, limitCancel := context.WithCancel(ctx)
master := &limiter{control: make(chan bool)}
batch := NewBatch(5, master)
promise := utils.PromiseCtx(limitCtx, batch.Start)
select {
case master.control <- true:
// we fed master and then cancel context
limitCancel()
case <-ctx.Done():
t.Fatal(ctx.Err())
}
select {
case err := <-promise:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal(ctx.Err())
}
}
func TestBatchLimiter(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*1)
defer cancel()
master := &limiter{control: make(chan bool, 2)}
for i := 0; i < 2; i++ {
master.control <- true
}
close(master.control)
batch := NewBatch(5, master)
promise := utils.PromiseCtx(ctx, batch.Start)
i, err := Drain(ctx, batch)
assert.NoError(t, err)
// we should take only 5 tick from master
assert.Equal(t, i, 10)
select {
case err := <-promise:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal(ctx.Err())
}
}
func TestContextCancelWhileControlSizeLimiter(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
sizeCtx, sizeCancel := context.WithCancel(ctx)
master := &limiter{control: make(chan struct{})}
size := NewSize(5, master)
promise := utils.PromiseCtx(sizeCtx, size.Start)
select {
case master.control <- struct{}{}:
// we fed master and then cancel context
sizeCancel()
case <-ctx.Done():
t.Fatal(ctx.Err())
}
select {
case err := <-promise:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal(ctx.Err())
}
}
func (bl *batch) Start(ctx context.Context) error {
defer close(bl.control)
masterCtx, cancelMaster := context.WithCancel(ctx)
masterPromise := utils.PromiseCtx(masterCtx, bl.master.Start)
loop:
for {
select {
case _, more := <-bl.master.Control():
if !more {
break loop
}
for i := 0; i < bl.batchSize; i++ {
select {
case bl.control <- struct{}{}:
case <-ctx.Done():
break loop
}
}
case <-ctx.Done():
break loop
}
}
cancelMaster()
err := <-masterPromise
return err
}
func (sl *sizeLimiter) Start(ctx context.Context) error {
defer close(sl.control)
masterCtx, cancelMaster := context.WithCancel(ctx)
masterPromise := utils.PromiseCtx(masterCtx, sl.master.Start)
loop:
for i := 0; i < sl.size; i++ {
select {
case v, more := <-sl.master.Control():
if !more {
break loop
}
select {
case sl.control <- v:
case <-ctx.Done():
break loop
}
case <-ctx.Done():
break loop
}
}
cancelMaster()
return <-masterPromise
}
func TestPeriodicLimiter(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*1)
defer cancel()
limiterCtx, cancelLimiter := context.WithCancel(ctx)
limiter := NewPeriodic(time.Millisecond * 2)
promise := utils.PromiseCtx(limiterCtx, limiter.Start)
ch := make(chan int)
go func() {
i, err := Drain(ctx, limiter)
if err != nil {
t.Fatal(err)
}
ch <- i
}()
time.Sleep(time.Millisecond * 7)
cancelLimiter()
select {
case i := <-ch:
// we should take only 4 ticks from ticker (1 in the beginning and 3 after 6 milliseconds)
assert.Equal(t, 4, i)
case <-ctx.Done():
t.Fatal(ctx.Err())
}
select {
case err := <-promise:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal(ctx.Err())
}
}
func TestLinearLimiter(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
defer cancel()
limiter := NewLinear(5, 6, 1)
promise := utils.PromiseCtx(ctx, limiter.Start)
ch := make(chan int, 100)
go func() {
i, err := Drain(ctx, limiter)
if err != nil {
t.Fatal(err)
}
ch <- i
}()
select {
case i := <-ch:
assert.Equal(t, 6, i)
case <-ctx.Done():
t.Fatal(ctx.Err())
}
select {
case err := <-promise:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal(ctx.Err())
}
}
func TestLinearLimiterFromConfig(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*1)
defer cancel()
lc := &config.Limiter{
LimiterType: "linear",
Parameters: map[string]interface{}{
"Period": 0.46,
"StartRps": 10.0,
"EndRps": 10.1,
},
}
limiter, err := NewLinearFromConfig(lc)
if err != nil {
t.Errorf("Got an error while creating linear limiter: %s", err)
}
if limiter == nil {
t.Errorf("Returned 'nil' with valid config")
}
switch tt := limiter.(type) {
case *linear:
default:
t.Errorf("Wrong limiter type returned (expected linear): %T", tt)
}
promise := utils.PromiseCtx(ctx, limiter.Start)
ch := make(chan int, 100)
go func() {
i, err := Drain(ctx, limiter)
if err != nil {
t.Fatal(err)
}
ch <- i
}()
select {
case i := <-ch:
assert.Equal(t, 5, i)
case <-ctx.Done():
t.Fatal(ctx.Err())
}
select {
case err := <-promise:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal(ctx.Err())
}
}
func TestHttpProvider(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
providerCtx, _ := context.WithCancel(ctx)
data, err := ioutil.ReadFile(httpTestFilename)
require.NoError(t, err)
ammoCh := make(chan Ammo, 128)
provider := &HttpProvider{
passes: 2,
ammoFile: bytes.NewReader(data),
sink: ammoCh,
BaseProvider: NewBaseProvider(
ammoCh,
HttpJSONDecode,
),
}
promise := utils.PromiseCtx(providerCtx, provider.Start)
ammos := Drain(ctx, provider)
require.Len(t, ammos, 25*2) // two passes
httpAmmo, casted := (ammos[2]).(*Http)
require.True(t, casted, "Ammo should have *Http type")
assert.Equal(t, "example.org", httpAmmo.Host)
assert.Equal(t, "/02", httpAmmo.Uri)
assert.Equal(t, "hello", httpAmmo.Tag)
assert.Equal(t, "GET", httpAmmo.Method)
assert.Len(t, httpAmmo.Headers, 4)
// TODO: add test for decoding error
select {
case err := <-promise:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal(ctx.Err())
}
}
func TestContextCancelInBatch(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
limitCtx, limitCancel := context.WithCancel(ctx)
limitCancel()
master := &limiter{control: make(chan bool, 10)}
batch := NewBatch(5, master)
promise := utils.PromiseCtx(limitCtx, batch.Start)
i, err := Drain(ctx, batch)
assert.NoError(t, err)
// we should take only 0 tick from master
assert.Equal(t, i, 0)
select {
case err := <-promise:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal(ctx.Err())
}
}
func TestContextCancelInSizeLimiter(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
sizeCtx, sizeCancel := context.WithCancel(ctx)
sizeCancel()
master := &limiter{control: make(chan struct{}, 10)}
size := NewSize(5, master)
promise := utils.PromiseCtx(sizeCtx, size.Start)
i, err := Drain(ctx, size)
assert.NoError(t, err)
// we should take only 0 tick from master
assert.Equal(t, i, 0)
select {
case err := <-promise:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal(ctx.Err())
}
}
func TestPeriodicLimiterBatchMaxCount(t *testing.T) {
lc := &config.Limiter{
LimiterType: "periodic",
Parameters: map[string]interface{}{
"Period": 0.46,
"BatchSize": 3.0,
"MaxCount": 5.0,
},
}
l, err := NewPeriodicFromConfig(lc)
if err != nil {
t.Errorf("Got an error while creating periodic limiter: %s", err)
}
if l == nil {
t.Errorf("Returned 'nil' with valid config")
}
switch tt := l.(type) {
case *batch:
default:
t.Errorf("Wrong limiter type returned (expected batchLimiter): %T", tt)
}
ctx, cancel := context.WithTimeout(context.Background(), time.Second*3)
defer cancel()
promise := utils.PromiseCtx(ctx, l.Start)
i, err := Drain(ctx, l)
assert.NoError(t, err)
// we should take only 0 tick from master
assert.Equal(t, i, 15)
select {
case err := <-promise:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal(ctx.Err())
}
}
func (e *Engine) Serve(ctx context.Context) error {
pools := make([]*UserPool, 0, len(e.cfg.Pools))
for _, upc := range e.cfg.Pools {
up, err := NewUserPoolFromConfig(&upc)
if err != nil {
log.Printf("Could not create user pool: %s", err)
continue
}
pools = append(pools, up)
}
promises := utils.Promises{}
for _, up := range pools {
promises = append(promises, utils.PromiseCtx(ctx, up.Start))
}
select {
case err := <-promises.All():
if err != nil {
return err
}
case <-ctx.Done():
}
log.Println("Done")
return nil
}
func TestContextCancelInUnlimited(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
limitCtx, limitCancel := context.WithCancel(ctx)
limitCancel()
lc := &config.Limiter{
LimiterType: "unlimited",
Parameters: nil,
}
unlimited, err := NewUnlimitedFromConfig(lc)
assert.NoError(t, err)
promise := utils.PromiseCtx(limitCtx, unlimited.Start)
_, err = Drain(ctx, unlimited)
assert.NoError(t, err)
select {
case err := <-promise:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal(ctx.Err())
}
}
func (up *UserPool) Start(ctx context.Context) error {
// userCtx will be canceled when all users finished their execution
userCtx, resultCancel := context.WithCancel(ctx)
userPromises := utils.Promises{}
utilsPromises := utils.Promises{
utils.PromiseCtx(ctx, up.ammunition.Start),
utils.PromiseCtx(userCtx, up.results.Start),
utils.PromiseCtx(ctx, up.startupLimiter.Start),
}
var sharedLimiter limiter.Limiter
if up.sharedSchedule {
var err error
sharedLimiter, err = GetLimiter(up.userLimiterConfig)
if err != nil {
return fmt.Errorf("could not make a user limiter from config due to %s", err)
}
// Starting shared limiter.
// This may cause spike load in the beginning of a test if it takes time
// to initialize a user, because we don't wait for them to initialize in
// case of shared limiter and there might be some ticks accumulated
utilsPromises = append(utilsPromises, utils.PromiseCtx(userCtx, sharedLimiter.Start))
}
for range up.startupLimiter.Control() {
var l limiter.Limiter
if up.sharedSchedule {
l = sharedLimiter
} else {
var err error
l, err = GetLimiter(up.userLimiterConfig)
if err != nil {
return fmt.Errorf("could not make a user limiter from config due to %s", err)
}
}
g, err := GetGun(up.gunConfig)
if err != nil {
return fmt.Errorf("could not make a gun from config due to %s", err)
}
u := &User{
Name: up.name,
Ammunition: up.ammunition,
Results: up.results,
Limiter: l,
Gun: g,
}
if !up.sharedSchedule {
utilsPromises = append(utilsPromises, utils.PromiseCtx(userCtx, l.Start))
}
userPromises = append(userPromises, utils.PromiseCtx(ctx, u.Run))
}
// FIXME: wrong logic here
log.Println("Started all users. Waiting for them")
err := <-userPromises.All()
resultCancel() // stop result listener when all users finished
err2 := utilsPromises.All()
if err2 != nil {
fmt.Printf("%v", err2)
}
return err
}