fix: race condition in ramp up

Author: rhamzeh

internal/requests/rampup.go

@@ -3,7 +3,9 @@ package requests
 import (
 	"context"
 	"fmt"
+	"math"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"golang.org/x/time/rate"
@@ -19,22 +21,38 @@ import (
 // - rampupPeriodDuration: The unit of each ramp-up period.
 // - maxInFlight: The maximum number of concurrent requests, used to protect the client and the server.
 // - requests: A slice of functions representing the API requests to be made.
-func RampUpAPIRequests( //nolint:gocognit,cyclop
+func RampUpAPIRequests( //nolint:cyclop
 	ctx context.Context,
 	minRPS, maxRPS, rampUpPeriod int, rampupPeriodDuration time.Duration, maxInFlight int,
 	requests []func() error,
 ) error {
-	rpsIncrement := float64(maxRPS-minRPS) / float64(rampUpPeriod)
-	limiter := rate.NewLimiter(rate.Limit(minRPS), 1)
-	semaphore := make(chan struct{}, maxInFlight)
+	var (
+		rpsIncrement = float64(maxRPS-minRPS) / float64(rampUpPeriod)
+		limiter      = rate.NewLimiter(rate.Limit(minRPS), 1)
+		semaphore    = make(chan struct{}, maxInFlight)
+		waitGroup    sync.WaitGroup
+		ticker       = time.NewTicker(rampupPeriodDuration)
+		requestIndex int32
+	)
 
-	var waitGroup sync.WaitGroup
-
-	ticker := time.NewTicker(rampupPeriodDuration)
 	defer ticker.Stop()
 
-	requestIndex := 0
-	requestsLen := len(requests)
+	if len(requests) > math.MaxInt32 {
+		return fmt.Errorf( //nolint:err113
+			"too many requests in ramp up: %d. max supported is %d", len(requests), math.MaxInt32,
+		)
+	}
+
+	requestsLen := int32(len(requests)) //nolint:gosec
+
+	worker := func(req func() error) {
+		defer waitGroup.Done()
+		defer func() { <-semaphore }()
+
+		if err := req(); err != nil {
+			fmt.Printf("Error: %v\n", err)
+		}
+	}
 
 	for step := 0; step <= rampUpPeriod; step++ {
 		select {
@@ -50,32 +68,19 @@ func RampUpAPIRequests( //nolint:gocognit,cyclop
 			}
 
 			for i := 0; i < int(limiter.Limit()); i++ { //nolint:intrange
-				if requestIndex >= requestsLen {
+				idx := atomic.AddInt32(&requestIndex, 1) - 1
+				if idx >= requestsLen {
 					waitGroup.Wait()
 
 					return nil
 				}
 
+				req := requests[idx]
 				semaphore <- struct{}{}
 
 				waitGroup.Add(1)
 
-				go func(req func() error) {
-					defer waitGroup.Done()
-					defer func() { <-semaphore }()
-
-					if req == nil {
-						fmt.Printf("Error: request function is nil, request %d out of %d\n", requestIndex, requestsLen)
-
-						return
-					}
-
-					if err := req(); err != nil {
-						fmt.Printf("Error: %v\n", err)
-					}
-				}(requests[requestIndex])
-
-				requestIndex++
+				go worker(req)
 			}
 
 			newRPS := rate.Limit(float64(minRPS) + rpsIncrement*float64(step))
@@ -99,26 +104,19 @@ func RampUpAPIRequests( //nolint:gocognit,cyclop
 			}
 
 			for i := 0; i < int(limiter.Limit()); i++ { //nolint:intrange
-				if requestIndex >= len(requests) {
+				idx := atomic.AddInt32(&requestIndex, 1) - 1
+				if idx >= requestsLen {
 					waitGroup.Wait()
 
 					return nil
 				}
 
+				req := requests[idx]
 				semaphore <- struct{}{}
 
 				waitGroup.Add(1)
 
-				go func(req func() error) {
-					defer waitGroup.Done()
-					defer func() { <-semaphore }()
-
-					if err := req(); err != nil {
-						fmt.Printf("Error: %v\n", err)
-					}
-				}(requests[requestIndex])
-
-				requestIndex++
+				go worker(req)
 			}
 		}
 	}

internal/requests/rampup_test.go

@@ -3,6 +3,7 @@ package requests_test
 import (
 	"context"
 	"errors"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -13,12 +14,17 @@ import (
 func TestRampUpAPIRequests_Success(t *testing.T) {
 	t.Parallel()
 
-	var callCount int32
+	var (
+		callCount int32
+		mutex     sync.Mutex
+	)
 
 	requestsList := make([]func() error, 5)
 	for i := range requestsList {
 		requestsList[i] = func() error {
+			mutex.Lock()
 			atomic.AddInt32(&callCount, 1)
+			mutex.Unlock()
 
 			return nil
 		}
@@ -29,7 +35,7 @@ func TestRampUpAPIRequests_Success(t *testing.T) {
 		t.Fatalf("expected no error, got %v, %v", err, callCount)
 	}
 
-	if atomic.LoadInt32(&callCount) != int32(len(requestsList)) { //nolint:gosec
+	if callCount != int32(len(requestsList)) { //nolint:gosec
 		t.Fatalf("expected %d calls, got %d", len(requestsList), callCount)
 	}
 }
@@ -40,12 +46,17 @@ func TestRampUpAPIRequests_RampUpRate(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
-	var callCount int32
+	var (
+		callCount int32
+		mutex     sync.Mutex
+	)
 
 	requestsList := make([]func() error, 10)
 	for i := range requestsList {
 		requestsList[i] = func() error {
+			mutex.Lock()
 			atomic.AddInt32(&callCount, 1)
+			mutex.Unlock()
 
 			return nil
 		}
@@ -72,12 +83,17 @@ func TestRampUpAPIRequests_ContextCancelled(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
 	defer cancel()
 
-	var callCount int32
+	var (
+		callCount int32
+		mutex     sync.Mutex
+	)
 
 	requestsList := make([]func() error, 100)
 	for i := range requestsList {
 		requestsList[i] = func() error {
+			mutex.Lock()
 			atomic.AddInt32(&callCount, 1)
+			mutex.Unlock()
 			time.Sleep(100 * time.Millisecond)
 
 			return nil
@@ -89,7 +105,7 @@ func TestRampUpAPIRequests_ContextCancelled(t *testing.T) {
 		t.Fatalf("expected error, got nil")
 	}
 
-	if atomic.LoadInt32(&callCount) == int32(len(requestsList)) { //nolint:gosec
+	if callCount == int32(len(requestsList)) { //nolint:gosec
 		t.Fatalf("expected fewer than %d calls, got %d", len(requestsList), callCount)
 	}
 }
@@ -100,16 +116,23 @@ func TestRampUpAPIRequests_RequestError(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 
-	var callCount int32
+	var (
+		callCount int32
+		mutex     sync.Mutex
+	)
 
 	requestsList := make([]func() error, 5)
 	for i := range requestsList {
 		requestsList[i] = func() error {
+			mutex.Lock()
 			atomic.AddInt32(&callCount, 1)
 
 			if callCount == 2 {
+				mutex.Unlock()
+
 				return errors.New("request error")
 			}
+			mutex.Unlock()
 
 			return nil
 		}
@@ -120,7 +143,7 @@ func TestRampUpAPIRequests_RequestError(t *testing.T) {
 		t.Fatalf("expected no error, got %v", err)
 	}
 
-	if atomic.LoadInt32(&callCount) != int32(len(requestsList)) { //nolint:gosec
+	if callCount != int32(len(requestsList)) { //nolint:gosec
 		t.Fatalf("expected %d calls, got %d", len(requestsList), callCount)
 	}
 }