Fix race condition with eager workflow start and close (#1191)

internal/internal_worker.go

@@ -951,6 +951,9 @@ func (aw *AggregatedWorker) Start() error {
 		if err := aw.workflowWorker.Start(); err != nil {
 			return err
 		}
+		if aw.client.eagerDispatcher != nil {
+			aw.client.eagerDispatcher.registerWorker(aw.workflowWorker)
+		}
 	}
 	if !util.IsInterfaceNil(aw.activityWorker) {
 		if err := aw.activityWorker.Start(); err != nil {
@@ -1568,9 +1571,6 @@ func NewAggregatedWorker(client *WorkflowClient, taskQueue string, options Worke
 		} else {
 			workflowWorker = newWorkflowWorker(client.workflowService, workerParams, nil, registry)
 		}
-		if client.eagerDispatcher != nil {
-			client.eagerDispatcher.registerWorker(workflowWorker)
-		}
 	}
 
 	// activity types.

internal/internal_worker_base.go

@@ -311,7 +311,7 @@ func (bw *baseWorker) runPoller() {
 }
 
 func (bw *baseWorker) tryReserveSlot() bool {
-	if !bw.isWorkerStarted || bw.isStop() {
+	if bw.isStop() {
 		return false
 	}
 	// Reserve a executor slot via a non-blocking attempt to take a poller

test/integration_test.go

@@ -1295,6 +1295,26 @@ func (ts *IntegrationTestSuite) TestBasicSession() {
 		ts.tracer.GetTrace("BasicSession"))
 }
 
+func (ts *IntegrationTestSuite) TestEagerWorkflowDispatchRaceWithWorkerStop() {
+	// Attempt to stop a worker while trying to schedule an eager workflow task
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		ctx := context.Background()
+		ctx, cancel := context.WithTimeout(ctx, ctxTimeout)
+		defer cancel()
+		_, err := ts.client.ExecuteWorkflow(ctx, ts.startWorkflowOptions("test-basic-session"), ts.workflows.SimplestWorkflow)
+		ts.NoError(err)
+		wg.Done()
+	}()
+	go func() {
+		ts.worker.Stop()
+		ts.workerStopped = true
+		wg.Done()
+	}()
+	wg.Wait()
+}
+
 func (ts *IntegrationTestSuite) TestSessionStateFailedWorkerFailed() {
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()