make slots configurable in load test

hatchet-dev · Jan 29, 2025 · 9534c0b · 9534c0b
1 parent 68250b6
commit 9534c0b
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 6 deletions.
diff --git a/examples/loadtest/cli/cli_e2e_test.go b/examples/loadtest/cli/cli_e2e_test.go
@@ -87,7 +87,7 @@ func TestLoadCLI(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			if err := do(tt.args.duration, tt.args.eventsPerSecond, tt.args.delay, tt.args.wait, tt.args.concurrency, tt.args.workerDelay); (err != nil) != tt.wantErr {
+			if err := do(tt.args.duration, tt.args.eventsPerSecond, tt.args.delay, tt.args.wait, tt.args.concurrency, tt.args.workerDelay, 100); (err != nil) != tt.wantErr {
 				t.Errorf("do() error = %v, wantErr %v", err, tt.wantErr)
 			}
 		})

diff --git a/examples/loadtest/cli/do.go b/examples/loadtest/cli/do.go
@@ -7,7 +7,7 @@ import (
 	"time"
 )
 
-func do(duration time.Duration, eventsPerSecond int, delay time.Duration, wait time.Duration, concurrency int, workerDelay time.Duration) error {
+func do(duration time.Duration, eventsPerSecond int, delay time.Duration, wait time.Duration, concurrency int, workerDelay time.Duration, slots int) error {
 	l.Info().Msgf("testing with duration=%s, eventsPerSecond=%d, delay=%s, wait=%s, concurrency=%d", duration, eventsPerSecond, delay, wait, concurrency)
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -28,7 +28,7 @@ func do(duration time.Duration, eventsPerSecond int, delay time.Duration, wait t
 			time.Sleep(workerDelay)
 		}
 		l.Info().Msg("starting worker now")
-		count, uniques := run(ctx, delay, durations, concurrency)
+		count, uniques := run(ctx, delay, durations, concurrency, slots)
 		ch <- count
 		ch <- uniques
 	}()

diff --git a/examples/loadtest/cli/main.go b/examples/loadtest/cli/main.go
@@ -21,6 +21,7 @@ func main() {
 	var delay time.Duration
 	var workerDelay time.Duration
 	var logLevel string
+	var slots int
 
 	var loadtest = &cobra.Command{
 		Use: "loadtest",
@@ -33,7 +34,7 @@ func main() {
 				"loadtest",
 			)
 
-			if err := do(duration, events, delay, wait, concurrency, workerDelay); err != nil {
+			if err := do(duration, events, delay, wait, concurrency, workerDelay, slots); err != nil {
 				log.Println(err)
 				panic("load test failed")
 			}
@@ -47,6 +48,7 @@ func main() {
 	loadtest.Flags().DurationVarP(&wait, "wait", "w", 10*time.Second, "wait specifies the total time to wait until events complete")
 	loadtest.Flags().DurationVarP(&workerDelay, "workerDelay", "p", 0*time.Second, "workerDelay specifies the time to wait before starting the worker")
 	loadtest.Flags().StringVarP(&logLevel, "level", "l", "info", "logLevel specifies the log level (debug, info, warn, error)")
+	loadtest.Flags().IntVarP(&slots, "slots", "s", 0, "slots specifies the number of slots to use in the worker")
 
 	cmd := &cobra.Command{Use: "app"}
 	cmd.AddCommand(loadtest)

diff --git a/examples/loadtest/cli/run.go b/examples/loadtest/cli/run.go
@@ -18,7 +18,7 @@ func getConcurrencyKey(ctx worker.HatchetContext) (string, error) {
 	return "my-key", nil
 }
 
-func run(ctx context.Context, delay time.Duration, executions chan<- time.Duration, concurrency int) (int64, int64) {
+func run(ctx context.Context, delay time.Duration, executions chan<- time.Duration, concurrency, slots int) (int64, int64) {
 	c, err := client.New(
 		client.WithLogLevel("warn"),
 	)
@@ -32,7 +32,7 @@ func run(ctx context.Context, delay time.Duration, executions chan<- time.Durati
 			c,
 		),
 		worker.WithLogLevel("warn"),
-		worker.WithMaxRuns(400),
+		worker.WithMaxRuns(slots),
 	)
 
 	if err != nil {