@@ -2,16 +2,28 @@ package loadgen
22
33import (
44 "context"
5+ "errors"
56 "fmt"
7+ "sync"
68 "time"
79
10+ "go.temporal.io/api/serviceerror"
811 "go.temporal.io/sdk/client"
912 "go.uber.org/zap"
1013)
1114
15+ // skipIterationErr is a sentinel error indicating that the iteration
16+ // should be skipped and not recorded as a completion or failure.
17+ var skipIterationErr = errors .New ("skip iteration" )
18+
1219type GenericExecutor struct {
1320 // Function to execute a single iteration of this scenario
1421 Execute func (context.Context , * Run ) error
22+
23+ // State management
24+ mu sync.Mutex
25+ state * ExecutorState
26+ workflowCompletionChecker * WorkflowCompletionChecker
1527}
1628
1729type genericRun struct {
@@ -24,13 +36,109 @@ type genericRun struct {
2436}
2537
2638func (g * GenericExecutor ) Run (ctx context.Context , info ScenarioInfo ) error {
39+ g .mu .Lock ()
40+ if g .state == nil {
41+ g .state = & ExecutorState {
42+ ExecutionID : info .ExecutionID ,
43+ }
44+ }
45+ if g .state .StartedAt .IsZero () {
46+ g .state .StartedAt = time .Now ()
47+ }
48+ g .mu .Unlock ()
49+
2750 r , err := g .newRun (info )
2851 if err != nil {
2952 return err
3053 }
3154 return r .Run (ctx )
3255}
3356
57+ func (g * GenericExecutor ) RecordCompletion () {
58+ g .mu .Lock ()
59+ defer g .mu .Unlock ()
60+
61+ if g .state != nil {
62+ g .state .CompletedIterations += 1
63+ g .state .LastCompletedAt = time .Now ()
64+ }
65+ }
66+
67+ func (g * GenericExecutor ) RecordError (err error ) {
68+ g .mu .Lock ()
69+ defer g .mu .Unlock ()
70+
71+ if g .state != nil && err != nil {
72+ g .state .IterationErrors = append (g .state .IterationErrors , err .Error ())
73+ }
74+ }
75+
76+ func (g * GenericExecutor ) VerifyRun (ctx context.Context , info ScenarioInfo ) []error {
77+ g .mu .Lock ()
78+ state := * g .state
79+ checker := g .workflowCompletionChecker
80+ g .mu .Unlock ()
81+
82+ if checker == nil {
83+ return nil
84+ }
85+ if err := checker .Verify (ctx , state ); err != nil {
86+ return []error {err }
87+ }
88+ return nil
89+ }
90+
91+ // EnableWorkflowCompletionCheck enables workflow completion verification for this executor.
92+ // It initializes a checker with the given timeout and registers the required search attributes.
93+ // The timeout specifies how long to wait for workflow completion verification (defaults to 30 seconds if zero).
94+ // The expectedWorkflowCount function, if provided, calculates the expected number of workflows from the ExecutorState.
95+ // If nil, defaults to using state.CompletedIterations.
96+ // Returns an error if search attribute registration fails.
97+ func (g * GenericExecutor ) EnableWorkflowCompletionCheck (ctx context.Context , info ScenarioInfo , timeout time.Duration , expectedWorkflowCount func (ExecutorState ) int ) error {
98+ checker , err := NewWorkflowCompletionChecker (ctx , info , timeout )
99+ if err != nil {
100+ return err
101+ }
102+
103+ if expectedWorkflowCount != nil {
104+ checker .SetExpectedWorkflowCount (expectedWorkflowCount )
105+ }
106+
107+ g .mu .Lock ()
108+ g .workflowCompletionChecker = checker
109+ g .mu .Unlock ()
110+
111+ return nil
112+ }
113+
114+ // GetState returns a copy of the current state
115+ func (g * GenericExecutor ) GetState () ExecutorState {
116+ g .mu .Lock ()
117+ defer g .mu .Unlock ()
118+
119+ if g .state == nil {
120+ return ExecutorState {}
121+ }
122+ return * g .state
123+ }
124+
125+ func (g * GenericExecutor ) Snapshot () any {
126+ return g .GetState ()
127+ }
128+
129+ func (g * GenericExecutor ) LoadState (loader func (any ) error ) error {
130+ var state ExecutorState
131+ if err := loader (& state ); err != nil {
132+ return err
133+ }
134+
135+ g .mu .Lock ()
136+ g .state = & state
137+ g .mu .Unlock ()
138+
139+ return nil
140+ }
141+
34142func (g * GenericExecutor ) newRun (info ScenarioInfo ) (* genericRun , error ) {
35143 info .Configuration .ApplyDefaults ()
36144 if err := info .Configuration .Validate (); err != nil {
@@ -83,7 +191,12 @@ func (g *genericRun) Run(ctx context.Context) error {
83191 case err := <- doneCh :
84192 currentlyRunning --
85193 if err != nil {
86- runErr = err
194+ if g .config .ContinueOnError {
195+ g .logger .Warnf ("Iteration failed but continuing due to --continue-on-error: %v" , err )
196+ g .executor .RecordError (err )
197+ } else {
198+ runErr = err
199+ }
87200 }
88201 case <- contextToWaitOn .Done ():
89202 }
@@ -130,25 +243,48 @@ func (g *genericRun) Run(ctx context.Context) error {
130243 defer func () {
131244 g .executeTimer .Record (time .Since (iterStart ))
132245
246+ // Check if this is the special "skip iteration" error
247+ isSkipIteration := errors .Is (err , skipIterationErr )
248+ if isSkipIteration {
249+ err = nil // Don't propagate this as an actual error
250+ }
251+
133252 select {
134253 case <- ctx .Done ():
135254 case doneCh <- err :
136- if err == nil && g .config .OnCompletion != nil {
137- g .config .OnCompletion (ctx , run )
255+ if err == nil && ! isSkipIteration {
256+ g .executor .RecordCompletion ()
257+ if g .config .OnCompletion != nil {
258+ g .config .OnCompletion (ctx , run )
259+ }
138260 }
139261 }
140262 }()
141263
142264 retryLoop:
143265 for {
144266 err = g .executor .Execute (ctx , run )
267+
268+ // Skip if workflow was already started.
269+ if err != nil {
270+ var alreadyStartedErr * serviceerror.WorkflowExecutionAlreadyStarted
271+ if errors .As (err , & alreadyStartedErr ) {
272+ g .logger .Debugf ("Workflow already started, skipping iteration %v" , run .Iteration )
273+ err = skipIterationErr
274+ break
275+ }
276+ }
277+
278+ // If defined, invoke user-defined error handler.
145279 if err != nil && g .config .HandleExecuteError != nil {
146280 err = g .config .HandleExecuteError (ctx , run , err )
147281 }
282+
148283 if err == nil {
149284 break
150285 }
151286
287+ // Attempt to retry.
152288 backoff , retry := run .ShouldRetry (err )
153289 if retry {
154290 err = fmt .Errorf ("iteration %v encountered error: %w" , run .Iteration , err )
0 commit comments