Fixed potential race-conditions; refactored def validation; do not perform skipped steps. (#23)

* Refactor workflow validation to centralize logic in `ValidateWorkflowDefinition`.

* Handle skipped, paused, and rolled-back steps in queue processing; correct error messages format in tests.

* Fix compensation retry count initialization from 1 to 0 in rollback processing.

* Handle skipped steps in queries and compensation retries; add max retries for nested onFailure steps.

* Fixed potential race-conditions and ensure null checks in store methods.

Author: rom8726

builder.go

@@ -555,7 +555,7 @@ func (builder *Builder) Build() (*WorkflowDefinition, error) {
 		},
 	}
 
-	if err := builder.validate(def); err != nil {
+	if err := ValidateWorkflowDefinition(def); err != nil {
 		return nil, err
 	}
 
@@ -568,30 +568,30 @@ func (builder *Builder) Build() (*WorkflowDefinition, error) {
 	return def, nil
 }
 
-func (builder *Builder) validate(def *WorkflowDefinition) error {
+func ValidateWorkflowDefinition(def *WorkflowDefinition) error {
 	for stepName, stepDef := range def.Definition.Steps {
 		if err := validateStepName(stepName); err != nil {
-			return fmt.Errorf("builder %q: %w", builder.name, err)
+			return fmt.Errorf("def %q: %w", def.Name, err)
 		}
 
 		for _, nextStep := range stepDef.Next {
 			if _, ok := def.Definition.Steps[nextStep]; !ok {
-				return fmt.Errorf("builder %q: step %q references unknown step: %q",
-					builder.name, stepName, nextStep)
+				return fmt.Errorf("def %q: step %q references unknown step: %q",
+					def.Name, stepName, nextStep)
 			}
 		}
 
 		if stepDef.OnFailure != "" {
 			if _, ok := def.Definition.Steps[stepDef.OnFailure]; !ok {
-				return fmt.Errorf("builder %q: step %q references unknown compensation step: %q",
-					builder.name, stepName, stepDef.OnFailure)
+				return fmt.Errorf("def %q: step %q references unknown compensation step: %q",
+					def.Name, stepName, stepDef.OnFailure)
 			}
 		}
 
 		if stepDef.Else != "" {
 			if _, ok := def.Definition.Steps[stepDef.Else]; !ok {
-				return fmt.Errorf("builder %q: step %q references unknown else step: %q",
-					builder.name, stepName, stepDef.Else)
+				return fmt.Errorf("def %q: step %q references unknown else step: %q",
+					def.Name, stepName, stepDef.Else)
 			}
 		}
 
@@ -601,8 +601,8 @@ func (builder *Builder) validate(def *WorkflowDefinition) error {
 				continue
 			}
 			if _, ok := def.Definition.Steps[parallelStep]; !ok {
-				return fmt.Errorf("builder %q: step %q references unknown parallel step: %q",
-					builder.name, stepName, parallelStep)
+				return fmt.Errorf("def %q: step %q references unknown parallel step: %q",
+					def.Name, stepName, parallelStep)
 			}
 		}
 
@@ -614,27 +614,27 @@ func (builder *Builder) validate(def *WorkflowDefinition) error {
 					continue
 				}
 				if _, ok := def.Definition.Steps[waitForStep]; !ok {
-					return fmt.Errorf("builder %q: join step %q references unknown step in waitFor: %q",
-						builder.name, stepName, waitForStep)
+					return fmt.Errorf("def %q: join step %q references unknown step in waitFor: %q",
+						def.Name, stepName, waitForStep)
 				}
 			}
 		}
 
 		if stepDef.Type == StepTypeTask && stepDef.Handler == "" {
-			return fmt.Errorf("builder %q: task step %q must have a handler", builder.name, stepName)
+			return fmt.Errorf("def %q: task step %q must have a handler", def.Name, stepName)
 		}
 	}
 
 	visited := make(map[string]bool)
-	err := builder.detectCycles(def.Definition.Start, def.Definition.Steps, visited, make(map[string]bool))
+	err := detectCycles(def.Definition.Start, def.Definition.Steps, visited, make(map[string]bool))
 	if err != nil {
-		return fmt.Errorf("builder %q: %w", builder.name, err)
+		return fmt.Errorf("def %q: %w", def.Name, err)
 	}
 
 	return nil
 }
 
-func (builder *Builder) detectCycles(
+func detectCycles(
 	current string,
 	steps map[string]*StepDefinition,
 	visited, recStack map[string]bool,
@@ -674,7 +674,7 @@ func (builder *Builder) detectCycles(
 		}
 
 		if !visited[next] {
-			if err := builder.detectCycles(next, steps, visited, recStack); err != nil {
+			if err := detectCycles(next, steps, visited, recStack); err != nil {
 				return err
 			}
 		} else if recStack[next] {
@@ -763,11 +763,11 @@ func (builder *Builder) traverseBranchForCondition(stepName string, visited map[
 	}
 
 	// Also check Else branch if exists
-	//if stepDef.Else != "" {
-	//	if result := builder.traverseBranchForCondition(stepDef.Else, visited); result != "" {
-	//		return result
-	//	}
-	//}
+	if stepDef.Else != "" {
+		if result := builder.traverseBranchForCondition(stepDef.Else, visited); result != "" {
+			return result
+		}
+	}
 
 	return ""
 }

engine.go

@@ -164,6 +164,10 @@ func (engine *Engine) Start(ctx context.Context, workflowID string, input json.R
 			return fmt.Errorf("get workflow definition: %w", err)
 		}
 
+		if err := engine.validateDefinition(def); err != nil {
+			return fmt.Errorf("invalid workflow definition: %w", err)
+		}
+
 		instance, err := engine.store.CreateInstance(ctx, workflowID, input)
 		if err != nil {
 			return fmt.Errorf("create instance: %w", err)
@@ -323,6 +327,17 @@ func (engine *Engine) ExecuteNext(ctx context.Context, workerID string) (empty b
 			}
 		}
 
+		// Do not execute skipped or paused steps
+		if step.Status == StepStatusSkipped ||
+			step.Status == StepStatusPaused ||
+			step.Status == StepStatusRolledBack {
+			if err := engine.store.RemoveFromQueue(ctx, step.ID); err != nil {
+				return fmt.Errorf("remove step from queue: %w", err)
+			}
+
+			return nil
+		}
+
 		// Check if this is a compensation
 		if step.Status == StepStatusCompensation {
 			// For distributed setup: if local engine doesn't have the compensation handler,
@@ -369,8 +384,6 @@ func (engine *Engine) ExecuteNext(ctx context.Context, workerID string) (empty b
 			}
 
 			return engine.executeCompensationStep(ctx, instance, step)
-		} else if step.Status == StepStatusRolledBack {
-			return nil
 		}
 
 		// Distributed handlers: if this is a task step and no local handler is registered,
@@ -919,7 +932,11 @@ func (engine *Engine) executeCompensationStep(ctx context.Context, instance *Wor
 			}
 
 			// Re-enqueue for retry
-			if err := engine.store.EnqueueStep(ctx, step.InstanceID, &step.ID, PriorityHigh, stepDef.Delay); err != nil {
+			retryDelay := CalculateRetryDelay(onFailureStep.RetryStrategy, onFailureStep.RetryDelay, newRetryCount)
+			if retryDelay == 0 {
+				retryDelay = onFailureStep.Delay
+			}
+			if err := engine.store.EnqueueStep(ctx, step.InstanceID, &step.ID, PriorityHigh, retryDelay); err != nil {
 				return fmt.Errorf("enqueue compensation retry: %w", err)
 			}
 
@@ -1659,11 +1676,6 @@ func (engine *Engine) notifyJoinStepsForStep(
 		return err
 	}
 
-	steps, err := engine.store.GetStepsByInstance(ctx, instanceID)
-	if err != nil {
-		return err
-	}
-
 	// Get Join step definition to check strategy
 	def, err := engine.store.GetWorkflowDefinition(ctx, instance.WorkflowID)
 	if err != nil {
@@ -1681,9 +1693,15 @@ func (engine *Engine) notifyJoinStepsForStep(
 		return fmt.Errorf("update join state for %s: %w", joinStepName, err)
 	}
 
+	var readySteps []WorkflowStep
 	// Additional check: don't consider join ready if there are still pending/running steps
 	if isReady {
-		hasPendingSteps := engine.hasPendingStepsInParallelBranches(ctx, instanceID, stepDef, steps)
+		readySteps, err = engine.store.GetStepsByInstance(ctx, instanceID)
+		if err != nil {
+			return err
+		}
+
+		hasPendingSteps := engine.hasPendingStepsInParallelBranches(ctx, instanceID, stepDef, readySteps)
 		if hasPendingSteps {
 			isReady = false
 			_, _ = engine.store.UpdateJoinState(ctx, instanceID, joinStepName, completedStepName, success)
@@ -1699,8 +1717,8 @@ func (engine *Engine) notifyJoinStepsForStep(
 
 	if isReady {
 		joinStepExists := false
-		for _, s := range steps {
-			if s.StepName == joinStepName {
+		for _, readyStep := range readySteps {
+			if readyStep.StepName == joinStepName {
 				joinStepExists = true
 
 				break
@@ -1709,9 +1727,9 @@ func (engine *Engine) notifyJoinStepsForStep(
 
 		if !joinStepExists {
 			var joinInput json.RawMessage
-			for _, s := range steps {
-				if s.StepName == completedStepName {
-					joinInput = s.Input
+			for _, readyStep := range readySteps {
+				if readyStep.StepName == completedStepName {
+					joinInput = readyStep.Input
 
 					break
 				}
@@ -1767,11 +1785,6 @@ func (engine *Engine) notifyJoinSteps(
 		return err
 	}
 
-	steps, err := engine.store.GetStepsByInstance(ctx, instanceID)
-	if err != nil {
-		return err
-	}
-
 	for stepName, stepDef := range def.Definition.Steps {
 		if stepDef.Type != StepTypeJoin {
 			continue
@@ -1815,8 +1828,14 @@ func (engine *Engine) notifyJoinSteps(
 
 		// Additional check: don't consider join ready if there are still pending/running steps
 		// in parallel branches that could affect the join result
+		var readySteps []WorkflowStep
 		if isReady {
-			hasPendingSteps := engine.hasPendingStepsInParallelBranches(ctx, instanceID, stepDef, steps)
+			readySteps, err = engine.store.GetStepsByInstance(ctx, instanceID)
+			if err != nil {
+				return err
+			}
+
+			hasPendingSteps := engine.hasPendingStepsInParallelBranches(ctx, instanceID, stepDef, readySteps)
 			if hasPendingSteps {
 				isReady = false
 				// Update the join state to reflect that it's not ready
@@ -1833,8 +1852,8 @@ func (engine *Engine) notifyJoinSteps(
 
 		if isReady {
 			joinStepExists := false
-			for _, s := range steps {
-				if s.StepName == stepName {
+			for _, readyStep := range readySteps {
+				if readyStep.StepName == stepName {
 					joinStepExists = true
 
 					break
@@ -1843,9 +1862,9 @@ func (engine *Engine) notifyJoinSteps(
 
 			if !joinStepExists {
 				var joinInput json.RawMessage
-				for _, s := range steps {
-					if s.StepName == completedStepName {
-						joinInput = s.Input
+				for _, readyStep := range readySteps {
+					if readyStep.StepName == completedStepName {
+						joinInput = readyStep.Input
 
 						break
 					}
@@ -2010,7 +2029,7 @@ func (engine *Engine) enqueueCompletedStepsForRollback(ctx context.Context, inst
 	// Mark each step as requiring compensation and enqueue for processing
 	for idx, step := range stepsToRollback {
 		// Update step status to compensation with retry count = 0
-		if err := engine.store.UpdateStepCompensationRetry(ctx, step.ID, 1, StepStatusCompensation); err != nil {
+		if err := engine.store.UpdateStepCompensationRetry(ctx, step.ID, 0, StepStatusCompensation); err != nil {
 			slog.Warn("[floxy] failed to mark step for compensation", "step_id", step.ID, "error", err)
 			continue
 		}
@@ -2215,28 +2234,7 @@ func (engine *Engine) validateDefinition(def *WorkflowDefinition) error {
 		return fmt.Errorf("start step not found: %s", def.Definition.Start)
 	}
 
-	for stepName, stepDef := range def.Definition.Steps {
-		for _, nextStep := range stepDef.Next {
-			if _, ok := def.Definition.Steps[nextStep]; !ok {
-				return fmt.Errorf("step %s references unknown step: %s", stepName, nextStep)
-			}
-		}
-
-		if stepDef.OnFailure != "" {
-			if _, ok := def.Definition.Steps[stepDef.OnFailure]; !ok {
-				return fmt.Errorf("step %s references unknown compensation step: %s",
-					stepName, stepDef.OnFailure)
-			}
-		}
-
-		for _, parallelStep := range stepDef.Parallel {
-			if _, ok := def.Definition.Steps[parallelStep]; !ok {
-				return fmt.Errorf("step %s references unknown parallel step: %s", stepName, parallelStep)
-			}
-		}
-	}
-
-	return nil
+	return ValidateWorkflowDefinition(def)
 }
 
 func (engine *Engine) rollbackToSavePointOrRoot(
@@ -2448,7 +2446,10 @@ func (engine *Engine) rollbackStep(ctx context.Context, step *WorkflowStep, def
 	}
 
 	// Enqueue compensation step for execution
-	retryDelay := CalculateRetryDelay(stepDef.RetryStrategy, stepDef.RetryDelay, newRetryCount)
+	retryDelay := CalculateRetryDelay(onFailureStep.RetryStrategy, onFailureStep.RetryDelay, newRetryCount)
+	if retryDelay == 0 {
+		retryDelay = onFailureStep.Delay
+	}
 	if err := engine.store.EnqueueStep(ctx, step.InstanceID, &step.ID, PriorityHigh, retryDelay); err != nil {
 		return fmt.Errorf("enqueue compensation step: %w", err)
 	}

engine_chaos_test.go

@@ -213,7 +213,7 @@ WHERE status = 'failed'`
 			nonRolledBackQuery := `
 				SELECT id, step_name, status
 				FROM workflows.workflow_steps
-				WHERE instance_id = $1 AND status != 'rolled_back'`
+				WHERE instance_id = $1 AND status != 'rolled_back' AND status != 'skipped'`
 			stepsRows, err := v.pool.Query(ctx, nonRolledBackQuery, instanceID)
 			if err != nil {
 				return fmt.Errorf("failed to query steps for instance %d: %w", instanceID, err)
@@ -245,7 +245,7 @@ WHERE status = 'failed'`
 				SELECT ws.id, ws.step_name, ws.status, ws.created_at::text
 				FROM workflows.workflow_steps ws
 				WHERE ws.instance_id = $1
-				  AND ws.status != 'rolled_back'
+				  AND ws.status != 'rolled_back' AND status != 'skipped'
 				  AND ws.created_at > (SELECT created_at FROM workflows.workflow_steps WHERE id = $2)
 				ORDER BY ws.created_at`
 			stepsRows, err := v.pool.Query(ctx, nonRolledBackQuery, instanceID, *savepointID)

engine_step_result_test.go

@@ -51,10 +51,6 @@ func Test_handleStepSuccess_SimpleNextFlow(t *testing.T) {
 	// 5) notifyJoinSteps path will need instance, def, steps list; provide empty steps list without joins
 	store.EXPECT().GetInstance(mock.Anything, instance.ID).Return(instance, nil)
 	store.EXPECT().GetWorkflowDefinition(mock.Anything, def.ID).Return(def, nil)
-	store.EXPECT().GetStepsByInstance(mock.Anything, instance.ID).Return([]WorkflowStep{
-		// include the just-finished step to provide input for potential join (not used here)
-		{ID: step.ID, InstanceID: instance.ID, StepName: step.StepName, StepType: step.StepType, Status: StepStatusCompleted, Input: json.RawMessage(`{"in":1}`)},
-	}, nil)
 	// 6) Enqueue next step B via enqueueNextSteps -> CreateStep + EnqueueStep
 	next := &WorkflowStep{
 		InstanceID: instance.ID,
@@ -142,9 +138,6 @@ func Test_handleStepFailure_DLQEnabled_NoRetry(t *testing.T) {
 	// 4) notifyJoinSteps (no joins) -> GetInstance, GetWorkflowDefinition, GetStepsByInstance
 	store.EXPECT().GetInstance(mock.Anything, instance.ID).Return(instance, nil)
 	store.EXPECT().GetWorkflowDefinition(mock.Anything, def.ID).Return(def, nil)
-	store.EXPECT().GetStepsByInstance(mock.Anything, instance.ID).Return([]WorkflowStep{
-		{ID: step.ID, InstanceID: instance.ID, StepName: step.StepName, StepType: step.StepType, Status: StepStatusPaused, Input: json.RawMessage(`{"in":1}`)},
-	}, nil)
 	// 5) Create DLQ record
 	store.EXPECT().CreateDeadLetterRecord(mock.Anything, mock.MatchedBy(func(rec *DeadLetterRecord) bool {
 		return rec != nil && rec.InstanceID == instance.ID && rec.StepID == step.ID && rec.WorkflowID == def.ID && rec.Reason != ""