Better shutdown pattern

Author: cschleiden

internal/worker/activity.go

@@ -52,7 +52,7 @@ func (aw *activityWorker) Start(ctx context.Context) error {
 		go aw.runPoll(ctx)
 	}
 
-	go aw.runDispatcher(ctx)
+	go aw.runDispatcher(context.Background())
 
 	return nil
 }
@@ -90,28 +90,25 @@ func (aw *activityWorker) runDispatcher(ctx context.Context) {
 		sem = make(chan struct{}, aw.options.MaxParallelActivityTasks)
 	}
 
-	for {
-		select {
-		case <-ctx.Done():
-			return
-		case task := <-aw.activityTaskQueue:
-			if sem != nil {
-				sem <- struct{}{}
-			}
+	for task := range aw.activityTaskQueue {
+		if sem != nil {
+			sem <- struct{}{}
+		}
 
-			aw.wg.Add(1)
-			go func() {
-				defer aw.wg.Done()
+		task := task
 
-				// Create new context to allow activities to complete when root context is canceled
-				taskCtx := context.Background()
-				aw.handleTask(taskCtx, task)
+		aw.wg.Add(1)
+		go func() {
+			defer aw.wg.Done()
 
-				if sem != nil {
-					<-sem
-				}
-			}()
-		}
+			// Create new context to allow activities to complete when root context is canceled
+			taskCtx := context.Background()
+			aw.handleTask(taskCtx, task)
+
+			if sem != nil {
+				<-sem
+			}
+		}()
 	}
 }
 

internal/worker/workflow.go

@@ -66,7 +66,7 @@ func (ww *workflowWorker) Start(ctx context.Context) error {
 		go ww.runPoll(ctx)
 	}
 
-	go ww.runDispatcher(ctx)
+	go ww.runDispatcher()
 
 	return nil
 }
@@ -99,33 +99,32 @@ func (ww *workflowWorker) runPoll(ctx context.Context) {
 	}
 }
 
-func (ww *workflowWorker) runDispatcher(ctx context.Context) {
+func (ww *workflowWorker) runDispatcher() {
 	var sem chan (struct{})
 
 	if ww.options.MaxParallelWorkflowTasks > 0 {
 		sem = make(chan struct{}, ww.options.MaxParallelWorkflowTasks)
 	}
 
-	for {
-		select {
-		case <-ctx.Done():
-			return
-		case t := <-ww.workflowTaskQueue:
-			if sem != nil {
-				sem <- struct{}{}
-			}
+	for t := range ww.workflowTaskQueue {
+		if sem != nil {
+			sem <- struct{}{}
+		}
 
-			ww.wg.Add(1)
-			go func() {
-				defer ww.wg.Done()
+		t := t
 
-				ww.handle(ctx, t)
+		ww.wg.Add(1)
+		go func() {
+			defer ww.wg.Done()
 
-				if sem != nil {
-					<-sem
-				}
-			}()
-		}
+			// Create new context to allow workflows to complete when root context is canceled
+			taskCtx := context.Background()
+			ww.handle(taskCtx, t)
+
+			if sem != nil {
+				<-sem
+			}
+		}()
 	}
 }