Retain original throughput starts across iteration retries

loadgen/kitchen_sink_executor.go

@@ -10,6 +10,9 @@ import (
 type KitchenSinkExecutor struct {
 	TestInput *kitchensink.TestInput
 
+	// RecoverOriginalStart retains the iteration request and acknowledged execution across retries.
+	RecoverOriginalStart bool
+
 	// Called once on start
 	PrepareTestInput func(context.Context, ScenarioInfo, *kitchensink.TestInput) error
 
@@ -27,7 +30,11 @@ func (k KitchenSinkExecutor) Run(ctx context.Context, info ScenarioInfo) error {
 	// Create generic executor and run it
 	ge := &GenericExecutor{
 		Execute: func(ctx context.Context, run *Run) error {
+			if k.RecoverOriginalStart && run.kitchenSinkExecution != nil {
+				return run.ExecuteKitchenSinkWorkflow(ctx, &run.kitchenSinkExecution.options)
+			}
 			options := run.DefaultKitchenSinkWorkflowOptions()
+			options.recoverOriginalStart = k.RecoverOriginalStart
 			testInputClone, ok := proto.Clone(k.TestInput).(*kitchensink.TestInput)
 			if !ok {
 				panic("failed to clone test input")

loadgen/kitchen_sink_start.go

@@ -0,0 +1,18 @@
+package loadgen
+
+import (
+	"github.com/golang/protobuf/proto"
+	"github.com/temporalio/omes/loadgen/kitchensink"
+)
+
+type kitchenSinkExecution struct {
+	options  KitchenSinkWorkflowOptions
+	executor *kitchensink.ClientActionsExecutor
+}
+
+func freezeKitchenSinkExecution(options KitchenSinkWorkflowOptions, executor *kitchensink.ClientActionsExecutor) *kitchenSinkExecution {
+	options.Params = proto.Clone(options.Params).(*kitchensink.TestInput)
+	executor.WorkflowInput = options.Params.GetWorkflowInput()
+	executor.RecoverOriginalStart = true
+	return &kitchenSinkExecution{options: options, executor: executor}
+}

loadgen/kitchen_sink_start_test.go

@@ -0,0 +1,55 @@
+package loadgen
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+	"github.com/temporalio/omes/loadgen/kitchensink"
+	"go.temporal.io/api/serviceerror"
+	"go.temporal.io/sdk/client"
+	"go.uber.org/zap"
+)
+
+type startInputsClient struct {
+	client.Client
+	ids    []string
+	inputs []int32
+}
+
+func (c *startInputsClient) ExecuteWorkflow(ctx context.Context, opts client.StartWorkflowOptions, _ interface{}, args ...interface{}) (client.WorkflowRun, error) {
+	c.ids = append(c.ids, opts.ID)
+	c.inputs = append(c.inputs, args[0].(*kitchensink.WorkflowInput).ExpectedSignalCount)
+	if len(c.ids) == 1 {
+		return nil, serviceerror.NewDeadlineExceeded("start response unavailable")
+	}
+	return completedOriginalRun{id: opts.ID}, nil
+}
+
+type completedOriginalRun struct {
+	client.WorkflowRun
+	id string
+}
+
+func (r completedOriginalRun) GetID() string                          { return r.id }
+func (r completedOriginalRun) GetRunID() string                       { return "original-run" }
+func (r completedOriginalRun) Get(context.Context, interface{}) error { return nil }
+
+func TestKitchenSinkOriginalStartFreezesRetryConstruction(t *testing.T) {
+	cl := &startInputsClient{}
+	generation := int32(0)
+	executor := KitchenSinkExecutor{TestInput: &kitchensink.TestInput{WorkflowInput: &kitchensink.WorkflowInput{}}, UpdateWorkflowOptions: func(ctx context.Context, r *Run, o *KitchenSinkWorkflowOptions) error {
+		generation++
+		o.StartOptions.ID = uuid.NewString()
+		o.Params.WorkflowInput.ExpectedSignalCount = generation
+		return nil
+	}}
+	executor.RecoverOriginalStart = true
+
+	info := ScenarioInfo{Client: cl, Namespace: "default", RunID: "frozen", ExecutionID: "same", Logger: zap.NewNop().Sugar(), MetricsHandler: client.MetricsNopHandler, Configuration: RunConfiguration{Iterations: 1, MaxConcurrent: 1, MaxIterationAttempts: 2, DoNotRegisterSearchAttributes: true}}
+	require.NoError(t, executor.Run(t.Context(), info))
+	require.Len(t, cl.ids, 2)
+	require.Equal(t, cl.ids[0], cl.ids[1])
+	require.Equal(t, []int32{1, 1}, cl.inputs)
+}

loadgen/kitchensink/client_action_executor.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/google/uuid"
 	enumspb "go.temporal.io/api/enums/v1"
+	"go.temporal.io/api/serviceerror"
 	"go.temporal.io/sdk/client"
 	"go.temporal.io/sdk/converter"
 	"go.temporal.io/sdk/workflow"
@@ -22,13 +23,40 @@ type ClientActionsExecutor struct {
 	WorkflowInput   *WorkflowInput
 	Handle          client.WorkflowRun
 	runID           string
+	// RecoverOriginalStart rejects duplicate starts and retains acknowledged run/update handles.
+	RecoverOriginalStart bool
+	originalStart        *originalStart
 }
 
 func (e *ClientActionsExecutor) Start(
 	ctx context.Context,
 	withStartAction *WithStartClientAction,
 ) error {
 	var err error
+	if e.RecoverOriginalStart {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+		if e.originalStart == nil {
+			if withStartAction.GetDoSignal() != nil {
+				return errors.New("original start recovery does not support Signal-with-Start")
+			}
+			e.freezeStart(withStartAction)
+		}
+		withStartAction = e.originalStart.action
+		if e.originalStart.update != nil {
+			if err := e.originalStart.update.Get(ctx, nil); err != nil {
+				return err
+			}
+		}
+		if e.Handle != nil {
+			e.runID = e.Handle.GetRunID()
+			return nil
+		}
+		if e.originalStart.withStartErr != nil {
+			return e.originalStart.withStartErr
+		}
+	}
 	if withStartAction == nil {
 		e.Handle, err = e.Client.ExecuteWorkflow(ctx, e.WorkflowOptions, e.WorkflowType, e.WorkflowInput)
 	} else if sig := withStartAction.GetDoSignal(); sig != nil {
@@ -38,6 +66,16 @@ func (e *ClientActionsExecutor) Start(
 	} else {
 		return fmt.Errorf("unsupported with_start_action: %v", withStartAction.String())
 	}
+	if err != nil && e.RecoverOriginalStart {
+		if withStartAction == nil {
+			var already *serviceerror.WorkflowExecutionAlreadyStarted
+			if errors.As(err, &already) && already.RunId != "" {
+				err = e.attachOriginalStart(ctx, already.RunId)
+			}
+		} else {
+			e.originalStart.withStartErr = err
+		}
+	}
 	if err != nil {
 		return fmt.Errorf("failed to start kitchen sink workflow: %w", err)
 	}
@@ -170,7 +208,11 @@ func (e *ClientActionsExecutor) executeUpdateAction(ctx context.Context, upd *Do
 	var handle client.WorkflowUpdateHandle
 	if upd.WithStart {
 		workflowOpts := e.WorkflowOptions
-		workflowOpts.WorkflowIDConflictPolicy = enumspb.WORKFLOW_ID_CONFLICT_POLICY_USE_EXISTING
+		if !e.RecoverOriginalStart {
+			workflowOpts.WorkflowIDConflictPolicy = enumspb.WORKFLOW_ID_CONFLICT_POLICY_USE_EXISTING
+		} else {
+			updateOpts.WaitForStage = client.WorkflowUpdateStageAccepted
+		}
 		op := e.Client.NewWithStartWorkflowOperation(workflowOpts, e.WorkflowType, e.WorkflowInput)
 		handle, err = e.Client.UpdateWithStartWorkflow(ctx, client.UpdateWithStartWorkflowOptions{
 			StartWorkflowOperation: op,
@@ -181,6 +223,10 @@ func (e *ClientActionsExecutor) executeUpdateAction(ctx context.Context, upd *Do
 	}
 
 	if err == nil {
+		if upd.WithStart && e.RecoverOriginalStart {
+			e.originalStart.update = handle
+			run = e.Client.GetWorkflow(ctx, handle.WorkflowID(), handle.RunID())
+		}
 		err = handle.Get(ctx, nil)
 		if upd.WithStart {
 			run = e.Client.GetWorkflow(ctx, handle.WorkflowID(), handle.RunID())

loadgen/kitchensink/original_start.go

@@ -0,0 +1,73 @@
+package kitchensink
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/google/uuid"
+	enumspb "go.temporal.io/api/enums/v1"
+	"go.temporal.io/sdk/client"
+	"go.temporal.io/sdk/converter"
+	"google.golang.org/protobuf/proto"
+)
+
+const startIntentMemo = "omes.original-start"
+
+type originalStart struct {
+	intent       string
+	action       *WithStartClientAction
+	update       client.WorkflowUpdateHandle
+	withStartErr error
+}
+
+func (e *ClientActionsExecutor) freezeStart(action *WithStartClientAction) {
+	e.originalStart = &originalStart{intent: uuid.NewString()}
+	memo := make(map[string]interface{}, len(e.WorkflowOptions.Memo)+1)
+	for k, v := range e.WorkflowOptions.Memo {
+		memo[k] = v
+	}
+	memo[startIntentMemo] = e.originalStart.intent
+	e.WorkflowOptions.Memo = memo
+	e.WorkflowOptions.WorkflowExecutionErrorWhenAlreadyStarted = true
+	e.WorkflowOptions.WorkflowIDReusePolicy = enumspb.WORKFLOW_ID_REUSE_POLICY_REJECT_DUPLICATE
+	e.WorkflowOptions.WorkflowIDConflictPolicy = enumspb.WORKFLOW_ID_CONFLICT_POLICY_FAIL
+	if upd := action.GetDoUpdate(); upd != nil && upd.UpdateId == "" {
+		upd.UpdateId = uuid.NewString()
+	}
+	if action != nil {
+		e.originalStart.action = proto.Clone(action).(*WithStartClientAction)
+	}
+}
+
+func (e *ClientActionsExecutor) attachOriginalStart(ctx context.Context, runID string) error {
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	history := e.Client.GetWorkflowHistory(ctx, e.WorkflowOptions.ID, runID, false, enumspb.HISTORY_EVENT_FILTER_TYPE_ALL_EVENT)
+	event, err := history.Next()
+	if err != nil {
+		return err
+	}
+	attrs := event.GetWorkflowExecutionStartedEventAttributes()
+	if attrs == nil {
+		return errors.New("original start history is missing its start event")
+	}
+	if attrs.GetMemo().GetFields()[startIntentMemo] == nil {
+		return errors.New("existing workflow has no original start identity")
+	}
+	var intent string
+	if err := converter.GetDefaultDataConverter().FromPayload(attrs.GetMemo().GetFields()[startIntentMemo], &intent); err != nil {
+		return fmt.Errorf("reading original start identity: %w", err)
+	}
+	input, err := converter.GetDefaultDataConverter().ToPayloads(e.WorkflowInput)
+	if err != nil {
+		return err
+	}
+	if intent != e.originalStart.intent || attrs.GetWorkflowType().GetName() != e.WorkflowType || attrs.GetTaskQueue().GetName() != e.WorkflowOptions.TaskQueue || !proto.Equal(attrs.GetInput(), input) {
+		return errors.New("existing workflow does not match the original start intent")
+	}
+	// Bind the history-verified run, never a later execution selected by workflow ID.
+	e.Handle = e.Client.GetWorkflow(ctx, e.WorkflowOptions.ID, runID)
+	return nil
+}