fix: resolve O(n^3) performance degradation in workflow tester (#377)

internal/workflowstate/workflowstate.go

@@ -57,6 +57,7 @@ type WfState struct {
 	instance        *core.WorkflowInstance
 	scheduleEventID int64
 	commands        []command.Command
+	commandIndex    map[int64]command.Command
 	pendingFutures  map[int64]*DecodingSettable
 	replaying       bool
 
@@ -75,6 +76,7 @@ func NewWorkflowState(instance *core.WorkflowInstance, logger *slog.Logger, trac
 	state := &WfState{
 		instance:        instance,
 		commands:        []command.Command{},
+		commandIndex:    make(map[int64]command.Command),
 		scheduleEventID: 1,
 		pendingFutures:  map[int64]*DecodingSettable{},
 
@@ -137,16 +139,11 @@ func (wf *WfState) Commands() []command.Command {
 
 func (wf *WfState) AddCommand(cmd command.Command) {
 	wf.commands = append(wf.commands, cmd)
+	wf.commandIndex[cmd.ID()] = cmd
 }
 
 func (wf *WfState) CommandByScheduleEventID(scheduleEventID int64) command.Command {
-	for _, c := range wf.commands {
-		if c.ID() == scheduleEventID {
-			return c
-		}
-	}
-
-	return nil
+	return wf.commandIndex[scheduleEventID]
 }
 
 func (wf *WfState) SetReplaying(replaying bool) {

tester/tester.go

@@ -74,6 +74,7 @@ type testWorkflow struct {
 	metadata      *metadata.WorkflowMetadata
 	history       []*history.Event
 	pendingEvents []*history.Event
+	executor      executor.WorkflowExecutor
 }
 
 type WorkflowTester[TResult any] interface {
@@ -357,30 +358,31 @@ func (wt *workflowTester[TResult]) Execute(ctx context.Context, args ...any) {
 			t := getNextWorkflowTask(tw.instance, tw.history, tw.pendingEvents)
 			tw.pendingEvents = tw.pendingEvents[:0]
 
-			// Execute task
-			e, err := executor.NewExecutor(
-				wt.logger,
-				wt.tracer,
-				wt.registry,
-				wt.converter,
-				wt.propagators,
-				&testHistoryProvider{tw.history},
-				tw.instance,
-				tw.metadata,
-				wt.clock,
-				wt.options.MaxHistorySize,
-			)
-			if err != nil {
-				panic(fmt.Errorf("could not create workflow executor: %v", err))
+			// Reuse cached executor or create a new one
+			if tw.executor == nil {
+				var err error
+				tw.executor, err = executor.NewExecutor(
+					wt.logger,
+					wt.tracer,
+					wt.registry,
+					wt.converter,
+					wt.propagators,
+					&testHistoryProvider{tw.history},
+					tw.instance,
+					tw.metadata,
+					wt.clock,
+					wt.options.MaxHistorySize,
+				)
+				if err != nil {
+					panic(fmt.Errorf("could not create workflow executor: %v", err))
+				}
 			}
 
-			result, err := e.ExecuteTask(ctx, t)
+			result, err := tw.executor.ExecuteTask(ctx, t)
 			if err != nil {
 				panic("Error while executing workflow" + err.Error())
 			}
 
-			e.Close()
-
 			// Add all executed events to history
 			tw.history = append(tw.history, result.Executed...)
 
@@ -397,6 +399,12 @@ func (wt *workflowTester[TResult]) Execute(ctx context.Context, args ...any) {
 						wt.workflowErr = a.Error
 					}
 
+					// Close cached executor for finished workflow
+					if tw.executor != nil {
+						tw.executor.Close()
+						tw.executor = nil
+					}
+
 				case history.EventType_TimerCanceled:
 					wt.cancelTimer(tw.instance, event)
 				}

tester/tester_bench_test.go

@@ -0,0 +1,49 @@
+package tester
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/cschleiden/go-workflows/workflow"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+// Reproduces https://github.com/cschleiden/go-workflows/issues/377
+func benchWF(ctx workflow.Context, n int) error {
+	wg := workflow.NewWaitGroup()
+	for range n {
+		wg.Add(1)
+		workflow.Go(ctx, func(ctx workflow.Context) {
+			defer wg.Done()
+			workflow.ExecuteActivity[any](ctx, workflow.DefaultActivityOptions, benchAct).Get(ctx)
+		})
+	}
+	wg.Wait(ctx)
+	return nil
+}
+
+func benchAct(ctx context.Context) error {
+	return nil
+}
+
+func BenchmarkWorkflowGo(b *testing.B) {
+	run := func(b *testing.B, n int) {
+		b.Run("n="+fmt.Sprint(n), func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				wft := NewWorkflowTester[any](benchWF)
+				wft.OnActivity(benchAct, mock.Anything).Return(nil)
+				wft.Execute(context.Background(), n)
+				require.True(b, wft.WorkflowFinished())
+				_, err := wft.WorkflowResult()
+				require.NoError(b, err)
+			}
+		})
+	}
+	run(b, 1)
+	run(b, 32)
+	run(b, 64)
+	run(b, 128)
+	run(b, 256)
+}