more exhaustive recovery test

maxdml · maxdml · commit ad463269ede3 · 2025-08-19T16:53:22.000-07:00
diff --git a/dbos/client_test.go b/dbos/client_test.go
@@ -235,7 +235,7 @@ func TestEnqueue(t *testing.T) {
 
 		// After first workflow completes, we should be able to enqueue with same deduplication ID
 		handle5, err := Enqueue[wfInput, string](clientCtx, queue.Name, "ServerWorkflow", wfInput{Input: "test-input"},
-			WithEnqueueWorkflowID(wfid2),   // Reuse the workflow ID that failed before
+			WithEnqueueWorkflowID(wfid2),        // Reuse the workflow ID that failed before
 			WithEnqueueDeduplicationID(dedupID), // Same deduplication ID as first workflow
 			WithEnqueueApplicationVersion(serverCtx.GetApplicationVersion()))
 		require.NoError(t, err, "failed to enqueue workflow with same dedup ID after completion")
diff --git a/dbos/serialization_test.go b/dbos/serialization_test.go
@@ -236,7 +236,7 @@ func TestSetEventSerialize(t *testing.T) {
 		assert.Equal(t, "user-defined-event-set", result)
 
 		// Retrieve the event to verify it was properly serialized and can be deserialized
-		retrievedEvent, err := GetEvent[UserDefinedEventData](executor, setHandle.GetWorkflowID(), "user-defined-key", 3 * time.Second)
+		retrievedEvent, err := GetEvent[UserDefinedEventData](executor, setHandle.GetWorkflowID(), "user-defined-key", 3*time.Second)
 		require.NoError(t, err)
 
 		// Verify the retrieved data matches what we set
@@ -273,7 +273,7 @@ func sendUserDefinedTypeWorkflow(ctx DBOSContext, destinationID string) (string,
 
 func recvUserDefinedTypeWorkflow(ctx DBOSContext, input string) (UserDefinedEventData, error) {
 	// Receive the user-defined type message
-	result, err := Recv[UserDefinedEventData](ctx, "user-defined-topic", 3 * time.Second)
+	result, err := Recv[UserDefinedEventData](ctx, "user-defined-topic", 3*time.Second)
 	return result, err
 }
 
diff --git a/dbos/workflows_test.go b/dbos/workflows_test.go
@@ -1009,57 +1009,148 @@ func TestWorkflowIdempotency(t *testing.T) {
 
 func TestWorkflowRecovery(t *testing.T) {
 	dbosCtx := setupDBOS(t, true, true)
-	RegisterWorkflow(dbosCtx, idempotencyWorkflowWithStep)
-	t.Run("RecoveryResumeWhereItLeftOff", func(t *testing.T) {
-		// Reset the global counter
-		idempotencyCounter = 0
 
-		// First execution - run the workflow once
-		input := "recovery-test"
-		idempotencyWorkflowWithStepEvent = NewEvent()
-		blockingStepStopEvent = NewEvent()
-		handle1, err := RunAsWorkflow(dbosCtx, idempotencyWorkflowWithStep, input)
-		require.NoError(t, err, "failed to execute workflow first time")
+	var (
+		recoveryCounters []int64
+		recoveryEvents   []*Event
+		blockingEvents   []*Event
+	)
+
+	recoveryWorkflow := func(dbosCtx DBOSContext, index int) (int64, error) {
+		// First step with custom name - increments the counter
+		_, err := RunAsStep(dbosCtx, func(ctx context.Context) (int64, error) {
+			recoveryCounters[index]++
+			return recoveryCounters[index], nil
+		}, WithStepName(fmt.Sprintf("IncrementStep-%d", index)))
+		if err != nil {
+			return 0, err
+		}
+
+		// Signal that first step is complete
+		recoveryEvents[index].Set()
+
+		// Second step with custom name - blocks until signaled
+		_, err = RunAsStep(dbosCtx, func(ctx context.Context) (string, error) {
+			blockingEvents[index].Wait()
+			return fmt.Sprintf("completed-%d", index), nil
+		}, WithStepName(fmt.Sprintf("BlockingStep-%d", index)))
+		if err != nil {
+			return 0, err
+		}
+
+		return recoveryCounters[index], nil
+	}
+
+	RegisterWorkflow(dbosCtx, recoveryWorkflow)
+
+	t.Run("WorkflowRecovery", func(t *testing.T) {
+		const numWorkflows = 5
+
+		// Initialize slices for multiple workflows
+		recoveryCounters = make([]int64, numWorkflows)
+		recoveryEvents = make([]*Event, numWorkflows)
+		blockingEvents = make([]*Event, numWorkflows)
+
+		// Create events for each workflow
+		for i := range numWorkflows {
+			recoveryEvents[i] = NewEvent()
+			blockingEvents[i] = NewEvent()
+		}
+
+		// Start all workflows
+		handles := make([]WorkflowHandle[int64], numWorkflows)
+		for i := range numWorkflows {
+			handle, err := RunAsWorkflow(dbosCtx, recoveryWorkflow, i, WithWorkflowID(fmt.Sprintf("recovery-test-%d", i)))
+			require.NoError(t, err, "failed to start workflow %d", i)
+			handles[i] = handle
+		}
+
+		// Wait for all first steps to complete
+		for i := range numWorkflows {
+			recoveryEvents[i].Wait()
+		}
+
+		// Verify step states before recovery
+		for i := range numWorkflows {
+			steps, err := dbosCtx.(*dbosContext).systemDB.getWorkflowSteps(dbosCtx, handles[i].GetWorkflowID())
+			require.NoError(t, err, "failed to get steps for workflow %d", i)
+			require.Len(t, steps, 1, "expected 1 completed step for workflow %d before recovery", i)
 
-		idempotencyWorkflowWithStepEvent.Wait() // Wait for the first step to complete. The second spins forever.
+			// Verify first step has custom name and completed
+			assert.Equal(t, fmt.Sprintf("IncrementStep-%d", i), steps[0].StepName, "workflow %d first step name mismatch", i)
+			assert.Equal(t, 0, steps[0].StepID, "workflow %d first step ID should be 0", i)
+			assert.NotNil(t, steps[0].Output, "workflow %d first step should have output", i)
+			assert.Nil(t, steps[0].Error, "workflow %d first step should not have error", i)
+		}
+
+		// Verify counters are all 1 (executed once)
+		for i := range numWorkflows {
+			require.Equal(t, int64(1), recoveryCounters[i], "workflow %d counter should be 1 before recovery", i)
+		}
 
-		// Run recovery for pending workflows with "local" executor
+		// Run recovery
 		recoveredHandles, err := recoverPendingWorkflows(dbosCtx.(*dbosContext), []string{"local"})
 		require.NoError(t, err, "failed to recover pending workflows")
+		require.Len(t, recoveredHandles, numWorkflows, "expected %d recovered handles, got %d", numWorkflows, len(recoveredHandles))
 
-		// Check that we have a single handle in the return list
-		require.Len(t, recoveredHandles, 1, "expected 1 recovered handle, got %d", len(recoveredHandles))
+		// Create a map for easy lookup of recovered handles
+		recoveredMap := make(map[string]WorkflowHandle[any])
+		for _, h := range recoveredHandles {
+			recoveredMap[h.GetWorkflowID()] = h
+		}
 
-		// Check that the workflow ID from the handle is the same as the first handle
-		recoveredHandle := recoveredHandles[0]
-		_, ok := recoveredHandle.(*workflowPollingHandle[any])
-		require.True(t, ok, "expected handle to be of type workflowPollingHandle, got %T", recoveredHandle)
-		require.Equal(t, handle1.GetWorkflowID(), recoveredHandle.GetWorkflowID())
+		// Verify all original workflows were recovered
+		for i := range numWorkflows {
+			originalID := handles[i].GetWorkflowID()
+			recoveredHandle, found := recoveredMap[originalID]
+			require.True(t, found, "workflow %d with ID %s not found in recovered handles", i, originalID)
+
+			_, ok := recoveredHandle.(*workflowPollingHandle[any])
+			require.True(t, ok, "recovered handle %d should be of type workflowPollingHandle, got %T", i, recoveredHandle)
+		}
 
-		idempotencyWorkflowWithStepEvent.Clear()
-		idempotencyWorkflowWithStepEvent.Wait()
+		// Verify first steps were NOT re-executed (counters should still be 1)
+		for i := range numWorkflows {
+			require.Equal(t, int64(1), recoveryCounters[i], "workflow %d counter should remain 1 after recovery (idempotent)", i)
+		}
 
-		// Check that the first step was *not* re-executed (idempotency counter is still 1)
-		require.Equal(t, int64(1), idempotencyCounter, "expected counter to remain 1 after recovery (idempotent)")
+		// Verify workflow attempts increased to 2
+		for i := range numWorkflows {
+			workflows, err := dbosCtx.(*dbosContext).systemDB.listWorkflows(context.Background(), listWorkflowsDBInput{
+				workflowIDs: []string{handles[i].GetWorkflowID()},
+			})
+			require.NoError(t, err, "failed to list workflow %d", i)
+			require.Len(t, workflows, 1, "expected 1 workflow entry for workflow %d", i)
+			assert.Equal(t, 2, workflows[0].Attempts, "workflow %d should have 2 attempts after recovery", i)
+		}
 
-		// Using ListWorkflows, retrieve the status of the workflow
-		workflows, err := dbosCtx.(*dbosContext).systemDB.listWorkflows(context.Background(), listWorkflowsDBInput{
-			workflowIDs: []string{handle1.GetWorkflowID()},
-		})
-		require.NoError(t, err, "failed to list workflows")
+		// Unblock all workflows and verify they complete
+		for i := range numWorkflows {
+			blockingEvents[i].Set()
+		}
 
-		require.Len(t, workflows, 1, "expected 1 workflow, got %d", len(workflows))
+		// Get results from all recovered workflows
+		for i := range numWorkflows {
+			recoveredHandle := recoveredMap[handles[i].GetWorkflowID()]
+			result, err := recoveredHandle.GetResult()
+			require.NoError(t, err, "failed to get result from recovered workflow %d", i)
 
-		workflow := workflows[0]
+			// Result should be the counter value (1)
+			require.Equal(t, int64(1), result, "workflow %d result should be 1", i)
+		}
 
-		// Ensure its number of attempts is 2
-		require.Equal(t, 2, workflow.Attempts)
+		// Final verification of step states
+		for i := range numWorkflows {
+			steps, err := dbosCtx.(*dbosContext).systemDB.getWorkflowSteps(dbosCtx, handles[i].GetWorkflowID())
+			require.NoError(t, err, "failed to get final steps for workflow %d", i)
+			require.Len(t, steps, 2, "expected 2 steps for workflow %d", i)
 
-		// unlock the workflow & wait for result
-		blockingStepStopEvent.Set() // This will allow the blocking step to complete
-		result, err := recoveredHandle.GetResult()
-		require.NoError(t, err, "failed to get result from recovered handle")
-		require.Equal(t, idempotencyCounter, result)
+			// Both steps should now be completed
+			assert.NotNil(t, steps[0].Output, "workflow %d first step should have output", i)
+			assert.NotNil(t, steps[1].Output, "workflow %d second step should have output", i)
+			assert.Nil(t, steps[0].Error, "workflow %d first step should not have error", i)
+			assert.Nil(t, steps[1].Error, "workflow %d second step should not have error", i)
+		}
 	})
 }
 
