Add rate.Limiter to RunConfiguration

loadgen/generic_executor.go

@@ -54,6 +54,9 @@ func (g *GenericExecutor) newRun(info ScenarioInfo) (*genericRun, error) {
 	if run.config.MaxConcurrent == 0 {
 		run.config.MaxConcurrent = g.DefaultConfiguration.MaxConcurrent
 	}
+	if run.config.Limiter == nil {
+		run.config.Limiter = g.DefaultConfiguration.Limiter
+	}
 	run.config.ApplyDefaults()
 	if run.config.Iterations > 0 && run.config.Duration > 0 {
 		return nil, fmt.Errorf("invalid scenario: iterations and duration are mutually exclusive")
@@ -104,8 +107,16 @@ func (g *genericRun) Run(ctx context.Context) error {
 		currentlyRunning++
 		run := g.info.NewRun(i + 1)
 		go func() {
-			startTime := time.Now()
-			err := g.executor.Execute(ctx, run)
+			var runStartTime time.Time
+			err := func() error {
+				if g.config.Limiter != nil {
+					if innerErr := g.config.Limiter.Wait(ctx); innerErr != nil {
+						return innerErr
+					}
+				}
+				runStartTime = time.Now()
+				return g.executor.Execute(ctx, run)
+			}()
 			// Only log/wrap/send to channel if context is not done
 			if ctx.Err() == nil {
 				if err != nil {
@@ -116,7 +127,7 @@ func (g *genericRun) Run(ctx context.Context) error {
 				case <-ctx.Done():
 				case doneCh <- err:
 					// Record/log here, not if it was cut short by context complete
-					g.executeTimer.Record(time.Since(startTime))
+					g.executeTimer.Record(time.Since(runStartTime))
 				}
 			}
 		}()

loadgen/scenario.go

@@ -3,17 +3,18 @@ package loadgen
 import (
 	"context"
 	"fmt"
-	"go.temporal.io/api/enums/v1"
-	"go.temporal.io/api/operatorservice/v1"
 	"path/filepath"
 	"runtime"
 	"strconv"
 	"strings"
 	"time"
 
 	"github.com/temporalio/omes/loadgen/kitchensink"
+	"go.temporal.io/api/enums/v1"
+	"go.temporal.io/api/operatorservice/v1"
 	"go.temporal.io/sdk/client"
 	"go.uber.org/zap"
+	"golang.org/x/time/rate"
 )
 
 type Scenario struct {
@@ -110,6 +111,9 @@ func (s *ScenarioInfo) ScenarioOptionInt(name string, defaultValue int) int {
 const DefaultIterations = 10
 const DefaultMaxConcurrent = 10
 
+// DefaultRateLimiter is unlimited
+var DefaultRateLimiter = rate.NewLimiter(rate.Inf, 0)
+
 type RunConfiguration struct {
 	// Number of iterations to run of this scenario (mutually exclusive with Duration).
 	Iterations int
@@ -119,6 +123,8 @@ type RunConfiguration struct {
 	// Maximum number of instances of the Execute method to run concurrently.
 	// Default is DefaultMaxConcurrent.
 	MaxConcurrent int
+	// Rate limiter to be Wait-ed before each iteration.
+	Limiter *rate.Limiter
 }
 
 func (r *RunConfiguration) ApplyDefaults() {
@@ -128,6 +134,9 @@ func (r *RunConfiguration) ApplyDefaults() {
 	if r.MaxConcurrent == 0 {
 		r.MaxConcurrent = DefaultMaxConcurrent
 	}
+	if r.Limiter == nil {
+		r.Limiter = DefaultRateLimiter
+	}
 }
 
 // Run represents an individual scenario run (many may be in a single instance (of possibly many) of a scenario).