Fix worker stop behavior

Author: cschleiden

backend/redis/redis_test.go

@@ -43,7 +43,6 @@ func createBackend() backend.Backend {
 		panic(err)
 	}
 
-	// Disable sticky workflow behavior for the test execution
 	b, err := NewRedisBackend(address, user, password, 0, WithBlockTimeout(time.Millisecond*2))
 	if err != nil {
 		panic(err)

backend/test/e2e.go

@@ -2,6 +2,7 @@ package test
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -15,10 +16,8 @@ import (
 
 func EndToEndBackendTest(t *testing.T, setup func() backend.Backend, teardown func(b backend.Backend)) {
 	tests := []struct {
-		name           string
-		workflow       interface{}
-		workflowInputs []interface{}
-		f              func(t *testing.T, ctx context.Context, c client.Client, w worker.Worker)
+		name string
+		f    func(t *testing.T, ctx context.Context, c client.Client, w worker.Worker)
 	}{
 		{
 			name: "SimpleWorkflow",
@@ -69,13 +68,18 @@ func EndToEndBackendTest(t *testing.T, setup func() backend.Backend, teardown fu
 		t.Run(tt.name, func(t *testing.T) {
 			b := setup()
 			ctx := context.Background()
+			ctx, cancel := context.WithCancel(ctx)
 
 			c := client.New(b)
 			w := worker.New(b, &worker.DefaultWorkerOptions)
 
 			tt.f(t, ctx, c, w)
 
-			w.Stop()
+			cancel()
+			if err := w.WaitForCompletion(); err != nil {
+				fmt.Println("Worker did not stop in time")
+				t.FailNow()
+			}
 
 			if teardown != nil {
 				teardown(b)

internal/worker/activity.go

@@ -16,7 +16,7 @@ import (
 
 type ActivityWorker interface {
 	Start(context.Context) error
-	Stop() error
+	WaitForCompletion() error
 }
 
 type activityWorker struct {
@@ -61,7 +61,7 @@ func (aw *activityWorker) Start(ctx context.Context) error {
 	return nil
 }
 
-func (aw *activityWorker) Stop() error {
+func (aw *activityWorker) WaitForCompletion() error {
 	aw.wg.Wait()
 
 	return nil

internal/worker/workflow.go

@@ -16,7 +16,7 @@ import (
 type WorkflowWorker interface {
 	Start(context.Context) error
 
-	Stop() error
+	WaitForCompletion() error
 }
 
 type workflowWorker struct {
@@ -64,7 +64,7 @@ func (ww *workflowWorker) Start(ctx context.Context) error {
 	return nil
 }
 
-func (ww *workflowWorker) Stop() error {
+func (ww *workflowWorker) WaitForCompletion() error {
 	ww.wg.Wait()
 
 	return nil

samples/simple/simple.go

@@ -34,7 +34,7 @@ func main() {
 
 	cancel()
 
-	if err := w.Stop(); err != nil {
+	if err := w.WaitForCompletion(); err != nil {
 		panic("could not stop worker" + err.Error())
 	}
 }

samples/subworkflow/subworkflow.go

@@ -36,7 +36,7 @@ func main() {
 
 	cancel()
 
-	if err := w.Stop(); err != nil {
+	if err := w.WaitForCompletion(); err != nil {
 		panic("could not stop worker" + err.Error())
 	}
 }

samples/timer/timer.go

@@ -32,7 +32,7 @@ func main() {
 	startWorkflow(ctx, c)
 
 	cancel()
-	w.Stop()
+	w.WaitForCompletion()
 }
 
 func startWorkflow(ctx context.Context, c client.Client) {

worker/worker.go

@@ -27,11 +27,14 @@ type Registry interface {
 type Worker interface {
 	Registry
 
-	// Start starts the worker
+	// Start starts the worker.
+	//
+	// To stop the worker, cancel the context passed to Start. To wait for completion of the active
+	// work items, call `WaitForCompletion`.
 	Start(ctx context.Context) error
 
-	// Stop stops the worker and waits for in-progress work to complete
-	Stop() error
+	// WaitForCompletion
+	WaitForCompletion() error
 }
 
 type worker struct {
@@ -83,12 +86,12 @@ func (w *worker) Start(ctx context.Context) error {
 	return nil
 }
 
-func (w *worker) Stop() error {
-	if err := w.workflowWorker.Stop(); err != nil {
+func (w *worker) WaitForCompletion() error {
+	if err := w.workflowWorker.WaitForCompletion(); err != nil {
 		return err
 	}
 
-	if err := w.activityWorker.Stop(); err != nil {
+	if err := w.activityWorker.WaitForCompletion(); err != nil {
 		return err
 	}