Change how workers wait for completion

Author: cschleiden

internal/worker/activity.go

@@ -25,7 +25,8 @@ type ActivityWorker struct {
 	activityTaskQueue    chan *task.Activity
 	activityTaskExecutor activity.Executor
 
-	wg *sync.WaitGroup
+	wg        sync.WaitGroup
+	pollersWg sync.WaitGroup
 
 	clock clock.Clock
 }
@@ -39,14 +40,14 @@ func NewActivityWorker(backend backend.Backend, registry *workflow.Registry, clo
 		activityTaskQueue:    make(chan *task.Activity),
 		activityTaskExecutor: activity.NewExecutor(backend.Logger(), backend.Tracer(), backend.Converter(), registry),
 
-		wg: &sync.WaitGroup{},
-
 		clock: clock,
 	}
 }
 
 func (aw *ActivityWorker) Start(ctx context.Context) error {
-	for i := 0; i <= aw.options.ActivityPollers; i++ {
+	aw.pollersWg.Add(aw.options.ActivityPollers)
+
+	for i := 0; i < aw.options.ActivityPollers; i++ {
 		go aw.runPoll(ctx)
 	}
 
@@ -56,9 +57,12 @@ func (aw *ActivityWorker) Start(ctx context.Context) error {
 }
 
 func (aw *ActivityWorker) WaitForCompletion() error {
-	close(aw.activityTaskQueue)
+	// Wait for task pollers to finish
+	aw.pollersWg.Wait()
 
+	// Wait for tasks to finish
 	aw.wg.Wait()
+	close(aw.activityTaskQueue)
 
 	return nil
 }

internal/worker/workflow.go

@@ -31,7 +31,8 @@ type WorkflowWorker struct {
 
 	logger log.Logger
 
-	wg *sync.WaitGroup
+	pollersWg sync.WaitGroup
+	wg        sync.WaitGroup
 }
 
 func NewWorkflowWorker(backend backend.Backend, registry *workflow.Registry, options *Options) *WorkflowWorker {
@@ -53,13 +54,13 @@ func NewWorkflowWorker(backend backend.Backend, registry *workflow.Registry, opt
 		cache: c,
 
 		logger: backend.Logger(),
-
-		wg: &sync.WaitGroup{},
 	}
 }
 
 func (ww *WorkflowWorker) Start(ctx context.Context) error {
-	for i := 0; i <= ww.options.WorkflowPollers; i++ {
+	ww.pollersWg.Add(ww.options.WorkflowPollers)
+
+	for i := 0; i < ww.options.WorkflowPollers; i++ {
 		go ww.runPoll(ctx)
 	}
 
@@ -69,14 +70,19 @@ func (ww *WorkflowWorker) Start(ctx context.Context) error {
 }
 
 func (ww *WorkflowWorker) WaitForCompletion() error {
-	close(ww.workflowTaskQueue)
+	// Wait for task pollers to finish
+	ww.pollersWg.Wait()
 
+	// Wait for tasks to finish
 	ww.wg.Wait()
+	close(ww.workflowTaskQueue)
 
 	return nil
 }
 
 func (ww *WorkflowWorker) runPoll(ctx context.Context) {
+	defer ww.pollersWg.Done()
+
 	for {
 		select {
 		case <-ctx.Done():