use a sync.Map for workflow registry

Author: maxdml

dbos/dbos.go

@@ -120,9 +120,8 @@ type dbosContext struct {
 	// Wait group for workflow goroutines
 	workflowsWg *sync.WaitGroup
 
-	// Workflow registry
-	workflowRegistry        map[string]workflowRegistryEntry
-	workflowRegMutex        *sync.RWMutex
+	// Workflow registry - read-mostly sync.Map since registration happens only before launch
+	workflowRegistry        sync.Map // map[string]workflowRegistryEntry
 	workflowCustomNametoFQN sync.Map // Maps fully qualified workflow names to custom names. Usefor when client enqueues a workflow by name because registry is indexed by FQN.
 
 	// Workflow scheduler
@@ -164,7 +163,6 @@ func WithValue(ctx DBOSContext, key, val any) DBOSContext {
 			systemDB:           dbosCtx.systemDB,
 			workflowsWg:        dbosCtx.workflowsWg,
 			workflowRegistry:   dbosCtx.workflowRegistry,
-			workflowRegMutex:   dbosCtx.workflowRegMutex,
 			applicationVersion: dbosCtx.applicationVersion,
 			executorID:         dbosCtx.executorID,
 			applicationID:      dbosCtx.applicationID,
@@ -187,7 +185,6 @@ func WithoutCancel(ctx DBOSContext) DBOSContext {
 			systemDB:           dbosCtx.systemDB,
 			workflowsWg:        dbosCtx.workflowsWg,
 			workflowRegistry:   dbosCtx.workflowRegistry,
-			workflowRegMutex:   dbosCtx.workflowRegMutex,
 			applicationVersion: dbosCtx.applicationVersion,
 			executorID:         dbosCtx.executorID,
 			applicationID:      dbosCtx.applicationID,
@@ -211,7 +208,6 @@ func WithTimeout(ctx DBOSContext, timeout time.Duration) (DBOSContext, context.C
 			systemDB:           dbosCtx.systemDB,
 			workflowsWg:        dbosCtx.workflowsWg,
 			workflowRegistry:   dbosCtx.workflowRegistry,
-			workflowRegMutex:   dbosCtx.workflowRegMutex,
 			applicationVersion: dbosCtx.applicationVersion,
 			executorID:         dbosCtx.executorID,
 			applicationID:      dbosCtx.applicationID,
@@ -262,11 +258,9 @@ func (c *dbosContext) GetApplicationID() string {
 func NewDBOSContext(inputConfig Config) (DBOSContext, error) {
 	ctx, cancelFunc := context.WithCancelCause(context.Background())
 	initExecutor := &dbosContext{
-		workflowsWg:      &sync.WaitGroup{},
-		ctx:              ctx,
-		ctxCancelFunc:    cancelFunc,
-		workflowRegistry: make(map[string]workflowRegistryEntry),
-		workflowRegMutex: &sync.RWMutex{},
+		workflowsWg:   &sync.WaitGroup{},
+		ctx:           ctx,
+		ctxCancelFunc: cancelFunc,
 	}
 
 	// Load and process the configuration

dbos/queue.go

@@ -210,11 +210,16 @@ func (qr *queueRunner) run(ctx *dbosContext) {
 					continue
 				}
 
-				registeredWorkflow, exists := ctx.workflowRegistry[wfName.(string)]
+				registeredWorkflowAny, exists := ctx.workflowRegistry.Load(wfName.(string))
 				if !exists {
 					ctx.logger.Error("workflow function not found in registry", "workflow_name", workflow.name)
 					continue
 				}
+				registeredWorkflow, ok := registeredWorkflowAny.(workflowRegistryEntry)
+				if !ok {
+					ctx.logger.Error("invalid workflow registry entry type", "workflow_name", workflow.name)
+					continue
+				}
 
 				// Deserialize input
 				var input any

dbos/recovery.go

@@ -43,11 +43,16 @@ func recoverPendingWorkflows(ctx *dbosContext, executorIDs []string) ([]Workflow
 			continue
 		}
 
-		registeredWorkflow, exists := ctx.workflowRegistry[wfName.(string)]
+		registeredWorkflowAny, exists := ctx.workflowRegistry.Load(wfName.(string))
 		if !exists {
 			ctx.logger.Error("Workflow function not found in registry", "workflow_id", workflow.ID, "name", workflow.Name)
 			continue
 		}
+		registeredWorkflow, ok := registeredWorkflowAny.(workflowRegistryEntry)
+		if !ok {
+			ctx.logger.Error("invalid workflow registry entry type", "workflow_id", workflow.ID, "name", workflow.Name)
+			continue
+		}
 
 		// Convert workflow parameters to options
 		opts := []WorkflowOption{

dbos/workflow.go

@@ -234,20 +234,17 @@ func registerWorkflow(ctx DBOSContext, workflowFQN string, fn WrappedWorkflowFun
 		panic("Cannot register workflow after DBOS has launched")
 	}
 
-	c.workflowRegMutex.Lock()
-	defer c.workflowRegMutex.Unlock()
-
-	if _, exists := c.workflowRegistry[workflowFQN]; exists {
-		c.logger.Error("workflow function already registered", "fqn", workflowFQN)
-		panic(newConflictingRegistrationError(workflowFQN))
-	}
-
-	// We must keep the registry indexed by FQN (because RunAsWorkflow uses reflection to find the function name and uses that to look it up in the registry)
-	c.workflowRegistry[workflowFQN] = workflowRegistryEntry{
+	// Check if workflow already exists and store atomically using LoadOrStore
+	entry := workflowRegistryEntry{
 		wrappedFunction: fn,
 		maxRetries:      maxRetries,
 		name:            customName,
 	}
+	
+	if _, exists := c.workflowRegistry.LoadOrStore(workflowFQN, entry); exists {
+		c.logger.Error("workflow function already registered", "fqn", workflowFQN)
+		panic(newConflictingRegistrationError(workflowFQN))
+	}
 
 	// We need to get a mapping from custom name to FQN for registry lookups that might not know the FQN (queue, recovery)
 	if len(customName) > 0 {
@@ -591,10 +588,14 @@ func (c *dbosContext) RunAsWorkflow(_ DBOSContext, fn WorkflowFunc, input any, o
 	}
 
 	// Lookup the registry for registration-time options
-	registeredWorkflow, exists := c.workflowRegistry[params.workflowName]
+	registeredWorkflowAny, exists := c.workflowRegistry.Load(params.workflowName)
 	if !exists {
 		return nil, newNonExistentWorkflowError(params.workflowName)
 	}
+	registeredWorkflow, ok := registeredWorkflowAny.(workflowRegistryEntry)
+	if !ok {
+		return nil, fmt.Errorf("invalid workflow registry entry type for workflow %s", params.workflowName)
+	}
 	if registeredWorkflow.maxRetries > 0 {
 		params.maxRetries = registeredWorkflow.maxRetries
 	}