Eliminate VU races, unify execution loop, remove cpu check loop (#1505)

skudasov · web-flow · commit cb02f6331d1b · 2024-12-19T11:34:25.000+01:00
diff --git a/.github/workflows/wasp-test.yml b/.github/workflows/wasp-test.yml
@@ -8,7 +8,7 @@ jobs:
     defaults:
       run:
         working-directory: wasp
-    runs-on: ubuntu-latest
+    runs-on: ubuntu22.04-16cores-64GB
     steps:
       - uses: actions/checkout@v3
       - uses: dorny/paths-filter@v3
diff --git a/wasp/stat.go b/wasp/stat.go
diff --git a/wasp/wasp.go b/wasp/wasp.go
@@ -217,7 +217,6 @@ type Stats struct {
 	CurrentSegment  atomic.Int64 `json:"current_schedule_segment"`
 	SamplesRecorded atomic.Int64 `json:"samples_recorded"`
 	SamplesSkipped  atomic.Int64 `json:"samples_skipped"`
-	RunStarted      atomic.Bool  `json:"runStarted"`
 	RunPaused       atomic.Bool  `json:"runPaused"`
 	RunStopped      atomic.Bool  `json:"runStopped"`
 	RunFailed       atomic.Bool  `json:"runFailed"`
@@ -246,6 +245,7 @@ type Generator struct {
 	Log                zerolog.Logger
 	labels             model.LabelSet
 	rl                 atomic.Pointer[ratelimit.Limiter]
+	rpsLoopOnce        *sync.Once
 	scheduleSegments   []*Segment
 	currentSegmentMu   *sync.Mutex
 	currentSegment     *Segment
@@ -320,6 +320,7 @@ func NewGenerator(cfg *Config) (*Generator, error) {
 		dataCancel:         dataCancel,
 		gun:                cfg.Gun,
 		vu:                 cfg.VU,
+		rpsLoopOnce:        &sync.Once{},
 		Responses:          NewResponses(rch),
 		ResponsesChan:      rch,
 		labels:             ls,
@@ -345,43 +346,26 @@ func NewGenerator(cfg *Config) (*Generator, error) {
 			return nil, err
 		}
 	}
-	CPUCheckLoop()
 	return g, nil
 }
 
-// runExecuteLoop initiates the generator's execution loop based on the configured load type.
-// It manages request pacing for RPS or handles virtual users for load testing scenarios.
-func (g *Generator) runExecuteLoop() {
-	g.currentSegment = g.scheduleSegments[0]
-	g.stats.LastSegment.Store(int64(len(g.scheduleSegments)))
-	switch g.Cfg.LoadType {
-	case RPS:
-		g.ResponsesWaitGroup.Add(1)
-		// we run pacedCall controlled by stats.CurrentRPS
-		go func() {
-			for {
-				select {
-				case <-g.ResponsesCtx.Done():
-					g.ResponsesWaitGroup.Done()
-					g.Log.Info().Msg("RPS generator has stopped")
-					return
-				default:
-					g.pacedCall()
-				}
+// runGunLoop runs the generator's Gun loop
+// It manages request pacing for RPS after the first segment is loaded.
+func (g *Generator) runGunLoop() {
+	g.ResponsesWaitGroup.Add(1)
+	// we run pacedCall controlled by stats.CurrentRPS
+	go func() {
+		for {
+			select {
+			case <-g.ResponsesCtx.Done():
+				g.ResponsesWaitGroup.Done()
+				g.Log.Info().Msg("RPS generator has stopped")
+				return
+			default:
+				g.pacedCall()
 			}
-		}()
-	case VU:
-		g.currentSegmentMu.Lock()
-		g.stats.CurrentVUs.Store(g.currentSegment.From)
-		g.currentSegmentMu.Unlock()
-		// we start all vus once
-		vus := g.stats.CurrentVUs.Load()
-		for i := 0; i < int(vus); i++ {
-			inst := g.vu.Clone(g)
-			g.runVU(inst)
-			g.vus = append(g.vus, inst)
 		}
-	}
+	}()
 }
 
 // runSetupWithTimeout executes the VirtualUser's setup within the configured timeout.
@@ -478,7 +462,6 @@ func (g *Generator) runVU(vu VirtualUser) {
 // It returns true when all segments have been handled, signaling the scheduler to terminate.
 func (g *Generator) processSegment() bool {
 	defer func() {
-		g.stats.RunStarted.Store(true)
 		g.Log.Info().
 			Int64("Segment", g.stats.CurrentSegment.Load()).
 			Int64("VUs", g.stats.CurrentVUs.Load()).
@@ -490,14 +473,18 @@ func (g *Generator) processSegment() bool {
 	}
 	g.currentSegmentMu.Lock()
 	g.currentSegment = g.scheduleSegments[g.stats.CurrentSegment.Load()]
-	g.currentSegment.StartTime = time.Now()
 	g.currentSegmentMu.Unlock()
 	g.stats.CurrentSegment.Add(1)
+	g.currentSegment.StartTime = time.Now()
 	switch g.Cfg.LoadType {
 	case RPS:
 		newRateLimit := ratelimit.New(int(g.currentSegment.From), ratelimit.Per(g.Cfg.RateLimitUnitDuration), ratelimit.WithoutSlack)
 		g.rl.Store(&newRateLimit)
 		g.stats.CurrentRPS.Store(g.currentSegment.From)
+		// start Gun loop once, in next segments we control it using g.rl ratelimiter
+		g.rpsLoopOnce.Do(func() {
+			g.runGunLoop()
+		})
 	case VU:
 		oldVUs := g.stats.CurrentVUs.Load()
 		newVUs := g.currentSegment.From
@@ -527,6 +514,8 @@ func (g *Generator) processSegment() bool {
 // runScheduleLoop initiates an asynchronous loop that processes scheduling segments and monitors for completion signals.
 // It enables the generator to handle load distribution seamlessly in the background.
 func (g *Generator) runScheduleLoop() {
+	g.currentSegment = g.scheduleSegments[0]
+	g.stats.LastSegment.Store(int64(len(g.scheduleSegments)))
 	go func() {
 		for {
 			select {
@@ -621,11 +610,7 @@ func (g *Generator) collectVUResults() {
 // handling timeouts and storing the response.
 // It ensures requests adhere to the generator's configuration and execution state.
 func (g *Generator) pacedCall() {
-	if !g.Stats().RunStarted.Load() {
-		return
-	}
-	l := *g.rl.Load()
-	l.Take()
+	(*g.rl.Load()).Take()
 	if g.stats.RunPaused.Load() {
 		return
 	}
@@ -668,7 +653,6 @@ func (g *Generator) Run(wait bool) (interface{}, bool) {
 		g.sendStatsToLoki()
 	}
 	g.runScheduleLoop()
-	g.runExecuteLoop()
 	g.collectVUResults()
 	if wait {
 		return g.Wait()
@@ -696,7 +680,6 @@ func (g *Generator) Stop() (interface{}, bool) {
 	if g.stats.RunStopped.Load() {
 		return nil, true
 	}
-	g.stats.RunStarted.Store(false)
 	g.stats.RunStopped.Store(true)
 	g.stats.RunFailed.Store(true)
 	g.Log.Warn().Msg("Graceful stop")
diff --git a/wasp/wasp_bench_test.go b/wasp/wasp_bench_test.go
@@ -21,7 +21,7 @@ func BenchmarkPacedCall(b *testing.B) {
 		Gun:               NewMockGun(&MockGunConfig{}),
 	})
 	require.NoError(b, err)
-	gen.runExecuteLoop()
+	gen.runGunLoop()
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		gen.pacedCall()
diff --git a/wasp/wasp_test.go b/wasp/wasp_test.go
@@ -449,13 +449,13 @@ func TestSmokeStaticRPSSchedulePrecision(t *testing.T) {
 	require.NoError(t, err)
 	_, failed := gen.Run(true)
 	require.Equal(t, false, failed)
-	require.GreaterOrEqual(t, gen.Stats().Success.Load(), int64(980))
+	require.GreaterOrEqual(t, gen.Stats().Success.Load(), int64(950))
 	require.LessOrEqual(t, gen.Stats().Success.Load(), int64(1010))
 	require.Equal(t, gen.Stats().Failed.Load(), int64(0))
 	require.Equal(t, gen.Stats().CallTimeout.Load(), int64(0))
 
 	okData, _, failResponses := convertResponsesData(gen)
-	require.GreaterOrEqual(t, len(okData), 980)
+	require.GreaterOrEqual(t, len(okData), 950)
 	require.LessOrEqual(t, len(okData), 1010)
 	require.Empty(t, failResponses)
 	require.Empty(t, gen.Errors())
@@ -475,14 +475,14 @@ func TestSmokeCustomUnitPrecision(t *testing.T) {
 	_, failed := gen.Run(true)
 	require.Equal(t, false, failed)
 	stats := gen.Stats()
-	require.GreaterOrEqual(t, stats.Success.Load(), int64(4970))
+	require.GreaterOrEqual(t, stats.Success.Load(), int64(4950))
 	require.LessOrEqual(t, stats.Success.Load(), int64(5010))
 	require.Equal(t, stats.Failed.Load(), int64(0))
 	require.Equal(t, stats.CallTimeout.Load(), int64(0))
 	require.Equal(t, stats.CurrentTimeUnit, gen.Cfg.RateLimitUnitDuration.Nanoseconds())
 
 	okData, _, failResponses := convertResponsesData(gen)
-	require.GreaterOrEqual(t, len(okData), 4970)
+	require.GreaterOrEqual(t, len(okData), 4950)
 	require.LessOrEqual(t, len(okData), 5010)
 	require.Empty(t, failResponses)
 	require.Empty(t, gen.Errors())
@@ -501,13 +501,13 @@ func TestSmokeStaticRPSScheduleIsNotBlocking(t *testing.T) {
 	require.NoError(t, err)
 	_, failed := gen.Run(true)
 	require.Equal(t, false, failed)
-	require.GreaterOrEqual(t, gen.Stats().Success.Load(), int64(980))
+	require.GreaterOrEqual(t, gen.Stats().Success.Load(), int64(950))
 	require.LessOrEqual(t, gen.Stats().Success.Load(), int64(1010))
 	require.Equal(t, gen.Stats().Failed.Load(), int64(0))
 	require.Equal(t, gen.Stats().CallTimeout.Load(), int64(0))
 
 	okData, _, failResponses := convertResponsesData(gen)
-	require.GreaterOrEqual(t, len(okData), 980)
+	require.GreaterOrEqual(t, len(okData), 950)
 	require.LessOrEqual(t, len(okData), 1010)
 	require.Empty(t, failResponses)
 	require.Empty(t, gen.Errors())
@@ -676,10 +676,10 @@ func TestSmokeVUsIncrease(t *testing.T) {
 
 	okData, okResponses, failResponses := convertResponsesData(gen)
 	require.GreaterOrEqual(t, okResponses[0].Duration, 50*time.Millisecond)
-	require.GreaterOrEqual(t, len(okResponses), 147)
-	require.GreaterOrEqual(t, len(okData), 147)
+	require.GreaterOrEqual(t, len(okResponses), 140)
+	require.GreaterOrEqual(t, len(okData), 140)
 	require.Equal(t, okResponses[0].Data.(string), "successCallData")
-	require.Equal(t, okResponses[147].Data.(string), "successCallData")
+	require.Equal(t, okResponses[140].Data.(string), "successCallData")
 	require.Empty(t, failResponses)
 	require.Empty(t, gen.Errors())
 }
@@ -706,10 +706,10 @@ func TestSmokeVUsDecrease(t *testing.T) {
 
 	okData, okResponses, failResponses := convertResponsesData(gen)
 	require.GreaterOrEqual(t, okResponses[0].Duration, 50*time.Millisecond)
-	require.GreaterOrEqual(t, len(okResponses), 147)
-	require.GreaterOrEqual(t, len(okData), 147)
+	require.GreaterOrEqual(t, len(okResponses), 140)
+	require.GreaterOrEqual(t, len(okData), 140)
 	require.Equal(t, okResponses[0].Data.(string), "successCallData")
-	require.Equal(t, okResponses[147].Data.(string), "successCallData")
+	require.Equal(t, okResponses[140].Data.(string), "successCallData")
 	require.Empty(t, failResponses)
 	require.Empty(t, gen.Errors())
 }
