Add failure-tolerant run mode (ContinueOnIterationFailure)

loadgen/generic_executor.go

@@ -3,6 +3,7 @@ package loadgen
 import (
 	"context"
 	"fmt"
+	"sync/atomic"
 	"time"
 
 	"go.temporal.io/sdk/client"
@@ -21,6 +22,10 @@ type genericRun struct {
 	logger   *zap.SugaredLogger
 	// Timer capturing E2E execution of each scenario run iteration.
 	executeTimer client.MetricsTimer
+	// Iteration outcome tallies, used for the end-of-run summary. failed counts
+	// only terminal failures tolerated via ContinueOnIterationFailure.
+	completed atomic.Int64
+	failed    atomic.Int64
 }
 
 func (g *GenericExecutor) Run(ctx context.Context, info ScenarioInfo) error {
@@ -130,11 +135,28 @@ func (g *genericRun) Run(ctx context.Context) error {
 			defer func() {
 				g.executeTimer.Record(time.Since(iterStart))
 
+				// A terminal iteration failure normally aborts the run (it is
+				// reported to the spawn loop via doneCh). When the run tolerates
+				// failures, swallow the error here so the loop keeps going; the
+				// failure is still surfaced via OnIterationFailure and the tally.
+				iterErr := err
+				if iterErr != nil && g.config.ContinueOnIterationFailure {
+					err = nil
+				}
+
 				select {
 				case <-ctx.Done():
 				case doneCh <- err:
-					if err == nil && g.config.OnCompletion != nil {
-						g.config.OnCompletion(ctx, run)
+					if iterErr == nil {
+						g.completed.Add(1)
+						if g.config.OnCompletion != nil {
+							g.config.OnCompletion(ctx, run)
+						}
+					} else {
+						g.failed.Add(1)
+						if g.config.OnIterationFailure != nil {
+							g.config.OnIterationFailure(ctx, run, iterErr)
+						}
 					}
 				}
 			}()
@@ -183,6 +205,11 @@ func (g *genericRun) Run(ctx context.Context) error {
 	if runErr != nil {
 		return fmt.Errorf("run finished with error after %v: %w", time.Since(startTime), runErr)
 	}
-	g.logger.Infof("Run completed in %v", time.Since(startTime))
+	if failed := g.failed.Load(); failed > 0 {
+		g.logger.Infof("Run completed in %v: %d iterations succeeded, %d failed (tolerated)",
+			time.Since(startTime), g.completed.Load(), failed)
+	} else {
+		g.logger.Infof("Run completed in %v", time.Since(startTime))
+	}
 	return nil
 }

loadgen/generic_executor_test.go

@@ -238,6 +238,45 @@ func TestExecutorRetries(t *testing.T) {
 	})
 }
 
+func TestRunContinueOnIterationFailure(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		var mu sync.Mutex
+		var completed, failed []int
+
+		err := execute(&GenericExecutor{
+			Execute: func(ctx context.Context, run *Run) error {
+				if run.Iteration%2 == 0 {
+					return errors.New("deliberate fail from test")
+				}
+				return nil
+			}},
+			RunConfiguration{
+				Iterations:                 6,
+				MaxConcurrent:              1,
+				ContinueOnIterationFailure: true,
+				OnCompletion: func(ctx context.Context, run *Run) {
+					mu.Lock()
+					defer mu.Unlock()
+					completed = append(completed, run.Iteration)
+				},
+				OnIterationFailure: func(ctx context.Context, run *Run, err error) {
+					mu.Lock()
+					defer mu.Unlock()
+					failed = append(failed, run.Iteration)
+				},
+			},
+		)
+
+		// Tolerated failures must not abort the run: every iteration runs and the
+		// outcomes are reported through the completion/failure hooks.
+		require.NoError(t, err)
+		mu.Lock()
+		defer mu.Unlock()
+		require.ElementsMatch(t, []int{1, 3, 5}, completed)
+		require.ElementsMatch(t, []int{2, 4, 6}, failed)
+	})
+}
+
 func TestExecutorRetriesLimit(t *testing.T) {
 	synctest.Test(t, func(t *testing.T) {
 		totalTracker := newIterationTracker()

loadgen/scenario.go

@@ -239,6 +239,17 @@ type RunConfiguration struct {
 	OnCompletion func(context.Context, *Run)
 	// HandleExecuteError, if set, is called when Execute returns an error, allowing transformation of errors.
 	HandleExecuteError func(context.Context, *Run, error) error
+	// ContinueOnIterationFailure, if set, keeps the run going after an iteration
+	// terminally fails (i.e. exhausts MaxIterationAttempts) instead of aborting on
+	// the first failure. The failed iteration is reported via OnIterationFailure
+	// and new iterations keep starting until the duration or iteration limit is
+	// reached. This is the mechanism for failure-tolerant runs; the policy for how
+	// many failures are acceptable belongs to the caller. Default is false.
+	ContinueOnIterationFailure bool
+	// OnIterationFailure, if set, is invoked when an iteration terminally fails
+	// (after exhausting retries). It is the hook callers use to tally failures
+	// when running with ContinueOnIterationFailure.
+	OnIterationFailure func(context.Context, *Run, error)
 }
 
 func (r *RunConfiguration) ApplyDefaults() {

scenarios/throughput_stress.go

@@ -51,6 +51,9 @@ const (
 type tpsState struct {
 	// CompletedIterations is the number of iteration that have been completed.
 	CompletedIterations int `json:"completedIterations"`
+	// FailedIterations is the number of iterations that terminally failed while
+	// the run tolerated failures (RunConfiguration.ContinueOnIterationFailure).
+	FailedIterations int `json:"failedIterations"`
 	// LastCompletedIterationAt is the time when the last iteration was completed. Helpful for debugging.
 	LastCompletedIterationAt time.Time `json:"lastCompletedIterationAt"`
 	// AccumulatedDuration is the total execution time across all runs (original + resumes).
@@ -215,6 +218,15 @@ func (t *tpsExecutor) Run(ctx context.Context, info loadgen.ScenarioInfo) error
 		info.Logger.Debugf("Completed iteration %d", run.Iteration)
 	}
 
+	// Tally terminal iteration failures into the state so the failed count is
+	// carried in the snapshot. This only fires when the run tolerates failures
+	// (RunConfiguration.ContinueOnIterationFailure); otherwise the first failure
+	// aborts the run before this is reached.
+	info.Configuration.OnIterationFailure = func(ctx context.Context, run *loadgen.Run, err error) {
+		t.updateStateOnIterationFailure()
+		info.Logger.Warnf("Iteration %d failed (tolerated): %v", run.Iteration, err)
+	}
+
 	// Start the scenario run.
 	//
 	// NOTE: When resuming, it can happen that there are no more iterations/time left to run more iterations.
@@ -338,8 +350,13 @@ func (t *tpsExecutor) Run(ctx context.Context, info loadgen.ScenarioInfo) error
 	}
 
 	// Post-scenario: ensure there are no failed or terminated workflows for this run.
-	if err := loadgen.VerifyNoFailedWorkflows(ctx, info, loadgen.OmesExecutionIDSearchAttribute, info.ExecutionID); err != nil {
-		tpsErrors = append(tpsErrors, err)
+	// When the run tolerates iteration failures, some workflows are expected to
+	// fail, so this zero-failure check is left to the caller's policy (which can
+	// read FailedIterations from the snapshot).
+	if !info.Configuration.ContinueOnIterationFailure {
+		if err := loadgen.VerifyNoFailedWorkflows(ctx, info, loadgen.OmesExecutionIDSearchAttribute, info.ExecutionID); err != nil {
+			tpsErrors = append(tpsErrors, err)
+		}
 	}
 
 	return errors.Join(tpsErrors...)
@@ -352,6 +369,12 @@ func (t *tpsExecutor) updateStateOnIterationCompletion() {
 	t.state.LastCompletedIterationAt = time.Now()
 }
 
+func (t *tpsExecutor) updateStateOnIterationFailure() {
+	t.lock.Lock()
+	defer t.lock.Unlock()
+	t.state.FailedIterations += 1
+}
+
 func (t *tpsExecutor) createActions(run *loadgen.Run) []*ActionSet {
 	return []*ActionSet{
 		{