wait for workflows to complete when shutting down the instance + fix bug in test

Author: maxdml

dbos/dbos.go

@@ -10,6 +10,7 @@ import (
 	"log/slog"
 	"net/url"
 	"os"
+	"sync"
 	"time"
 
 	"github.com/robfig/cron/v3"
@@ -93,6 +94,7 @@ type executor struct {
 	applicationVersion    string
 	applicationID         string
 	executorID            string
+	workflowsWg           *sync.WaitGroup
 }
 
 func Initialize(inputConfig Config) error {
@@ -101,7 +103,9 @@ func Initialize(inputConfig Config) error {
 		return newInitializationError("DBOS already initialized")
 	}
 
-	initExecutor := &executor{}
+	initExecutor := &executor{
+		workflowsWg: &sync.WaitGroup{},
+	}
 
 	// Load & process the configuration
 	config, err := processConfig(&inputConfig)
@@ -204,6 +208,7 @@ func Shutdown() {
 	}
 
 	// XXX is there a way to ensure all workflows goroutine are done before closing?
+	dbos.workflowsWg.Wait()
 
 	// Cancel the context to stop the queue runner
 	if dbos.queueRunnerCancelFunc != nil {

dbos/workflow.go

@@ -517,7 +517,9 @@ func runAsWorkflow[P any, R any](ctx context.Context, fn WorkflowFunc[P, R], inp
 
 	// Run the function in a goroutine
 	augmentUserContext := context.WithValue(ctx, workflowStateKey, wfState)
+	dbos.workflowsWg.Add(1)
 	go func() {
+		defer dbos.workflowsWg.Done()
 		result, err := fn(augmentUserContext, input)
 		status := WorkflowStatusSuccess
 		if err != nil {

dbos/workflows_test.go

@@ -788,6 +788,8 @@ func TestWorkflowDeadLetterQueue(t *testing.T) {
 			t.Fatalf("expected DeadLetterQueueError, got %v", dbosErr.Code)
 		}
 
+		// Unlock the workflow to allow it to complete
+		deadLetterQueueEvent.Set()
 		/*
 				// TODO: test resume when implemented
 				resumedHandle, err := ...