From 2a7d7f9725230845c4dad4d109dba055742eec7b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ivan=20Miri=C4=87?= <ivan@loadimpact.com>
Date: Tue, 11 Feb 2020 18:32:28 +0100
Subject: [PATCH] Fix variable-looping-vus VU wobble during ramp-down bug

Part 1 of #1296
---
 lib/executor/variable_looping_vus.go      |   2 +-
 lib/executor/variable_looping_vus_test.go | 100 ++++++++++++++++++++--
 2 files changed, 92 insertions(+), 10 deletions(-)

diff --git a/lib/executor/variable_looping_vus.go b/lib/executor/variable_looping_vus.go
index 63ae4dd1ad1..f763cdc0fbd 100644
--- a/lib/executor/variable_looping_vus.go
+++ b/lib/executor/variable_looping_vus.go
@@ -590,7 +590,7 @@ func (vlv VariableLoopingVUs) Run(ctx context.Context, out chan<- stats.SampleCo
 				}
 				handleNewScheduledVUs(step.PlannedVUs)
 			case step := <-gracefulLimitEvents:
-				if step.PlannedVUs > currentScheduledVUs {
+				if step.PlannedVUs > currentMaxAllowedVUs {
 					// Handle the case where a value is read from the
 					// gracefulLimitEvents channel before rawStepEvents
 					handleNewScheduledVUs(step.PlannedVUs)
diff --git a/lib/executor/variable_looping_vus_test.go b/lib/executor/variable_looping_vus_test.go
index 5d2c110d23b..45242f0ed4a 100644
--- a/lib/executor/variable_looping_vus_test.go
+++ b/lib/executor/variable_looping_vus_test.go
@@ -22,6 +22,7 @@ package executor
 
 import (
 	"context"
+	"fmt"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -35,10 +36,13 @@ import (
 	"github.com/loadimpact/k6/lib/types"
 )
 
-func getTestVariableLoopingVUsConfig() VariableLoopingVUsConfig {
-	return VariableLoopingVUsConfig{
-		BaseConfig: BaseConfig{GracefulStop: types.NullDurationFrom(0)},
-		StartVUs:   null.IntFrom(5),
+func TestVariableLoopingVUsRun(t *testing.T) {
+	t.Parallel()
+
+	config := VariableLoopingVUsConfig{
+		BaseConfig:       BaseConfig{GracefulStop: types.NullDurationFrom(0)},
+		GracefulRampDown: types.NullDurationFrom(0),
+		StartVUs:         null.IntFrom(5),
 		Stages: []Stage{
 			{
 				Duration: types.NullDurationFrom(1 * time.Second),
@@ -53,16 +57,12 @@ func getTestVariableLoopingVUsConfig() VariableLoopingVUsConfig {
 				Target:   null.IntFrom(3),
 			},
 		},
-		GracefulRampDown: types.NullDurationFrom(0),
 	}
-}
 
-func TestVariableLoopingVUsRun(t *testing.T) {
-	t.Parallel()
 	var iterCount int64
 	es := lib.NewExecutionState(lib.Options{}, 10, 50)
 	var ctx, cancel, executor, _ = setupExecutor(
-		t, getTestVariableLoopingVUsConfig(), es,
+		t, config, es,
 		simpleRunner(func(ctx context.Context) error {
 			time.Sleep(200 * time.Millisecond)
 			atomic.AddInt64(&iterCount, 1)
@@ -94,3 +94,85 @@ func TestVariableLoopingVUsRun(t *testing.T) {
 	assert.Equal(t, []int64{5, 3, 0}, result)
 	assert.Equal(t, int64(40), iterCount)
 }
+
+// Ensure there's no wobble of VUs during graceful ramp-down, without segments.
+// See https://github.com/loadimpact/k6/issues/1296
+func TestVariableLoopingVUsRampDownNoWobble(t *testing.T) {
+	t.Parallel()
+
+	config := VariableLoopingVUsConfig{
+		BaseConfig:       BaseConfig{GracefulStop: types.NullDurationFrom(0)},
+		GracefulRampDown: types.NullDurationFrom(1 * time.Second),
+		StartVUs:         null.IntFrom(0),
+		Stages: []Stage{
+			{
+				Duration: types.NullDurationFrom(3 * time.Second),
+				Target:   null.IntFrom(10),
+			},
+			{
+				Duration: types.NullDurationFrom(2 * time.Second),
+				Target:   null.IntFrom(0),
+			},
+		},
+	}
+
+	es := lib.NewExecutionState(lib.Options{}, 10, 50)
+	var ctx, cancel, executor, _ = setupExecutor(
+		t, config, es,
+		simpleRunner(func(ctx context.Context) error {
+			time.Sleep(1 * time.Second)
+			return nil
+		}),
+	)
+	defer cancel()
+
+	var (
+		wg     sync.WaitGroup
+		result []int64
+		m      sync.Mutex
+	)
+
+	sampleActiveVUs := func(delay time.Duration) {
+		time.Sleep(delay)
+		m.Lock()
+		result = append(result, es.GetCurrentlyActiveVUsCount())
+		m.Unlock()
+	}
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		sampleActiveVUs(100 * time.Millisecond)
+		sampleActiveVUs(3 * time.Second)
+		time.AfterFunc(2*time.Second, func() {
+			sampleActiveVUs(0)
+		})
+		time.Sleep(1 * time.Second)
+		// Sample ramp-down at a higher frequency
+		for i := 0; i < 15; i++ {
+			sampleActiveVUs(100 * time.Millisecond)
+		}
+	}()
+
+	err := executor.Run(ctx, nil)
+
+	wg.Wait()
+	require.NoError(t, err)
+	assert.Equal(t, int64(0), result[0])
+	assert.Equal(t, int64(10), result[1])
+	assert.Equal(t, int64(0), result[len(result)-1])
+
+	var curr int64
+	last := result[2]
+	// Check all ramp-down samples
+	for i := 3; i < len(result[2:]); i++ {
+		curr = result[i]
+		// Detect ramp-ups, missteps (e.g. 7 -> 4), but ignore pauses
+		if curr > last || (curr != last && curr != last-1) {
+			assert.FailNow(t,
+				fmt.Sprintf("ramping down wobble bug - "+
+					"current: %d, previous: %d\nVU samples: %v", curr, last, result))
+		}
+		last = curr
+	}
+}