assign stepID

Author: af-md

dbos/workflow.go

@@ -987,18 +987,18 @@ type StepFunc func(ctx context.Context) (any, error)
 // Step represents a type-safe step function with a specific output type R.
 type Step[R any] func(ctx context.Context) (R, error)
 
-// stepOptions holds the configuration for step execution using functional options pattern.
-type stepOptions struct {
-	maxRetries    int           // Maximum number of retry attempts (0 = no retries)
-	backoffFactor float64       // Exponential backoff multiplier between retries (default: 2.0)
-	baseInterval  time.Duration // Initial delay between retries (default: 100ms)
-	maxInterval   time.Duration // Maximum delay between retries (default: 5s)
-	stepName      string        // Custom name for the step (defaults to function name)
-	stepID        int
+// StepOptions holds the configuration for step execution using functional options pattern.
+type StepOptions struct {
+	maxRetries         int           // Maximum number of retry attempts (0 = no retries)
+	backoffFactor      float64       // Exponential backoff multiplier between retries (default: 2.0)
+	baseInterval       time.Duration // Initial delay between retries (default: 100ms)
+	maxInterval        time.Duration // Maximum delay between retries (default: 5s)
+	stepName           string        // Custom name for the step (defaults to function name)
+	preGeneratedStepID *int          // Pre generated stepID in case we want to run the function in a Go routine
 }
 
 // setDefaults applies default values to stepOptions
-func (opts *stepOptions) setDefaults() {
+func (opts *StepOptions) setDefaults() {
 	if opts.backoffFactor == 0 {
 		opts.backoffFactor = _DEFAULT_STEP_BACKOFF_FACTOR
 	}
@@ -1011,12 +1011,12 @@ func (opts *stepOptions) setDefaults() {
 }
 
 // StepOption is a functional option for configuring step execution parameters.
-type StepOption func(*stepOptions)
+type StepOption func(*StepOptions)
 
 // WithStepName sets a custom name for the step. If the step name has already been set
 // by a previous call to WithStepName, this option will be ignored
 func WithStepName(name string) StepOption {
-	return func(opts *stepOptions) {
+	return func(opts *StepOptions) {
 		if opts.stepName == "" {
 			opts.stepName = name
 		}
@@ -1026,7 +1026,7 @@ func WithStepName(name string) StepOption {
 // WithStepMaxRetries sets the maximum number of retry attempts for the step.
 // A value of 0 means no retries (default behavior).
 func WithStepMaxRetries(maxRetries int) StepOption {
-	return func(opts *stepOptions) {
+	return func(opts *StepOptions) {
 		opts.maxRetries = maxRetries
 	}
 }
@@ -1035,31 +1035,31 @@ func WithStepMaxRetries(maxRetries int) StepOption {
 // The delay between retries is calculated as: BaseInterval * (BackoffFactor^(retry-1))
 // Default value is 2.0.
 func WithBackoffFactor(factor float64) StepOption {
-	return func(opts *stepOptions) {
+	return func(opts *StepOptions) {
 		opts.backoffFactor = factor
 	}
 }
 
 // WithBaseInterval sets the initial delay between retries.
 // Default value is 100ms.
 func WithBaseInterval(interval time.Duration) StepOption {
-	return func(opts *stepOptions) {
+	return func(opts *StepOptions) {
 		opts.baseInterval = interval
 	}
 }
 
 // WithMaxInterval sets the maximum delay between retries.
 // Default value is 5s.
 func WithMaxInterval(interval time.Duration) StepOption {
-	return func(opts *stepOptions) {
+	return func(opts *StepOptions) {
 		opts.maxInterval = interval
 	}
 }
 
 
 func WithNextStepID(stepID int) StepOption {
-	return func(opts *stepOptions) {
-		opts.stepID = stepID
+	return func(opts *StepOptions) {
+		opts.preGeneratedStepID = &stepID
 	}
 }
 
@@ -1142,7 +1142,7 @@ func RunAsStep[R any](ctx DBOSContext, fn Step[R], opts ...StepOption) (R, error
 
 func (c *dbosContext) RunAsStep(_ DBOSContext, fn StepFunc, opts ...StepOption) (any, error) {
 	// Process functional options
-	stepOpts := &stepOptions{}
+	stepOpts := &StepOptions{}
 	for _, opt := range opts {
 		opt(stepOpts)
 	}
@@ -1164,18 +1164,21 @@ func (c *dbosContext) RunAsStep(_ DBOSContext, fn StepFunc, opts ...StepOption)
 		return fn(c)
 	}
 
+	// Get stepID if it has been pre generated
+	var stepID int
+	if stepOpts.preGeneratedStepID != nil {
+		stepID = *stepOpts.preGeneratedStepID
+	} else {
+		stepID = wfState.nextStepID() // crucially, this increments the step ID on the *workflow* state
+	}
+
 	// Setup step state
 	stepState := workflowState{
 		workflowID:   wfState.workflowID,
-		stepID:       wfState.nextStepID(), // crucially, this increments the step ID on the *workflow* state
+		stepID:       stepID,
 		isWithinStep: true,
 	}
 
-	// this logic needs to be looked at
-	if stepOpts.stepID >= 0 {
-		stepState.stepID = stepOpts.stepID
-	}
-
 	// Uncancellable context for DBOS operations
 	uncancellableCtx := WithoutCancel(c)
 
@@ -1278,7 +1281,7 @@ func Go[R any](ctx DBOSContext, fn Step[R], opts ...StepOption) (R, error) {
 	if !ok || wfState == nil {
 		return *new(R), newStepExecutionError("", stepName, "workflow state not found in context: are you running this step within a workflow?")
 	}
-	stepID := wfState.NextStepID()
+	stepID := wfState.nextStepID()
 	opts = append(opts, WithNextStepID(stepID))
 
 	// Type-erase the function