fix leaky goroutines in matching

cadence-workflow · Dec 19, 2023 · 3741401 · 3741401
1 parent e869984
commit 3741401
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/service/matching/taskListManager_test.go b/service/matching/taskListManager_test.go
@@ -31,6 +31,7 @@ import (
 	"github.com/pborman/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
 	"github.com/uber/cadence/common/cache"
 	"github.com/uber/cadence/common/cluster"
 	"github.com/uber/cadence/common/dynamicconfig"

diff --git a/service/matching/taskReader.go b/service/matching/taskReader.go
@@ -26,6 +26,7 @@ import (
 	"context"
 	"errors"
 	"runtime"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -80,6 +81,9 @@ type (
 		dispatchTask             func(context.Context, *InternalTask) error
 		getIsolationGroupForTask func(context.Context, *persistence.TaskInfo) (string, error)
 		ratePerSecond            func() float64
+
+		// stopWg is used to wait for all dispatchers to stop.
+		stopWg sync.WaitGroup
 	}
 )
 
@@ -137,6 +141,7 @@ func (tr *taskReader) Stop() {
 				tag.Error(err))
 		}
 		tr.taskGC.RunNow(tr.taskAckManager.GetAckLevel())
+		tr.stopWg.Wait()
 	}
 }
 
@@ -149,6 +154,8 @@ func (tr *taskReader) Signal() {
 }
 
 func (tr *taskReader) dispatchBufferedTasks(isolationGroup string) {
+	tr.stopWg.Add(1)
+	defer tr.stopWg.Done()
 dispatchLoop:
 	for {
 		select {
@@ -168,6 +175,9 @@ dispatchLoop:
 }
 
 func (tr *taskReader) getTasksPump() {
+	tr.stopWg.Add(1)
+	defer tr.stopWg.Done()
+
 	updateAckTimer := time.NewTimer(tr.config.UpdateAckInterval())
 	defer updateAckTimer.Stop()
 getTasksPumpLoop: