Fix potential go-routine leak when retrying local activity (#918)

* Add stopCh to local activity tunnel to fix go-routine leak

Author: yycptt

internal/internal_task_handlers.go

@@ -976,7 +976,10 @@ func (w *workflowExecutionContextImpl) retryLocalActivity(lar *localActivityResu
 			}
 
 			lar.task.attempt++
-			w.laTunnel.sendTask(lar.task)
+
+			if !w.laTunnel.sendTask(lar.task) {
+				lar.task.attempt--
+			}
 		})
 		return true
 	}
@@ -1055,13 +1058,18 @@ func (w *workflowExecutionContextImpl) CompleteDecisionTask(workflowTask *workfl
 	if w.hasPendingLocalActivityWork() && w.laTunnel != nil {
 		if len(eventHandler.unstartedLaTasks) > 0 {
 			// start new local activity tasks
+			unstartedLaTasks := make(map[string]struct{})
 			for activityID := range eventHandler.unstartedLaTasks {
 				task := eventHandler.pendingLaTasks[activityID]
 				task.wc = w
 				task.workflowTask = workflowTask
-				w.laTunnel.sendTask(task)
+				if !w.laTunnel.sendTask(task) {
+					unstartedLaTasks[activityID] = struct{}{}
+					task.wc = nil
+					task.workflowTask = nil
+				}
 			}
-			eventHandler.unstartedLaTasks = make(map[string]struct{})
+			eventHandler.unstartedLaTasks = unstartedLaTasks
 		}
 		// cannot complete decision task as there are pending local activities
 		if waitLocalActivities {

internal/internal_task_handlers_test.go

@@ -1106,18 +1106,18 @@ func (t *TaskHandlersTestSuite) TestLocalActivityRetry_DecisionHeartbeatFail() {
 	}
 
 	task := createWorkflowTask(testEvents, 0, "RetryLocalActivityWorkflow")
+	stopCh := make(chan struct{})
 	params := workerExecutionParameters{
-		TaskList: testWorkflowTaskTasklist,
-		Identity: "test-id-1",
-		Logger:   t.logger,
-		Tracer:   opentracing.NoopTracer{},
+		TaskList:          testWorkflowTaskTasklist,
+		Identity:          "test-id-1",
+		Logger:            t.logger,
+		Tracer:            opentracing.NoopTracer{},
+		WorkerStopChannel: stopCh,
 	}
+	defer close(stopCh)
 
 	taskHandler := newWorkflowTaskHandler(testDomain, params, nil, getHostEnvironment())
-	laTunnel := &localActivityTunnel{
-		taskCh:   make(chan *localActivityTask, 1000),
-		resultCh: make(chan interface{}),
-	}
+	laTunnel := newLocalActivityTunnel(params.WorkerStopChannel)
 	taskHandlerImpl, ok := taskHandler.(*workflowTaskHandlerImpl)
 	t.True(ok)
 	taskHandlerImpl.laTunnel = laTunnel

internal/internal_task_pollers.go

@@ -132,20 +132,34 @@ type (
 	localActivityTunnel struct {
 		taskCh   chan *localActivityTask
 		resultCh chan interface{}
+		stopCh   <-chan struct{}
 	}
 )
 
-func (lat *localActivityTunnel) getTask(stopC <-chan struct{}) *localActivityTask {
+func newLocalActivityTunnel(stopCh <-chan struct{}) *localActivityTunnel {
+	return &localActivityTunnel{
+		taskCh:   make(chan *localActivityTask, 1000),
+		resultCh: make(chan interface{}),
+		stopCh:   stopCh,
+	}
+}
+
+func (lat *localActivityTunnel) getTask() *localActivityTask {
 	select {
 	case task := <-lat.taskCh:
 		return task
-	case <-stopC:
+	case <-lat.stopCh:
 		return nil
 	}
 }
 
-func (lat *localActivityTunnel) sendTask(task *localActivityTask) {
-	lat.taskCh <- task
+func (lat *localActivityTunnel) sendTask(task *localActivityTask) bool {
+	select {
+	case lat.taskCh <- task:
+		return true
+	case <-lat.stopCh:
+		return false
+	}
 }
 
 func isClientSideError(err error) bool {
@@ -414,7 +428,7 @@ func newLocalActivityPoller(params workerExecutionParameters, laTunnel *localAct
 }
 
 func (latp *localActivityTaskPoller) PollTask() (interface{}, error) {
-	return latp.laTunnel.getTask(latp.shutdownC), nil
+	return latp.laTunnel.getTask(), nil
 }
 
 func (latp *localActivityTaskPoller) ProcessTask(task interface{}) error {

internal/internal_worker.go

@@ -308,10 +308,7 @@ func newWorkflowTaskWorkerInternal(
 	)
 
 	// laTunnel is the glue that hookup 3 parts
-	laTunnel := &localActivityTunnel{
-		taskCh:   make(chan *localActivityTask, 1000),
-		resultCh: make(chan interface{}),
-	}
+	laTunnel := newLocalActivityTunnel(params.WorkerStopChannel)
 
 	// 1) workflow handler will send local activity task to laTunnel
 	if handlerImpl, ok := taskHandler.(*workflowTaskHandlerImpl); ok {

test/integration_test.go

@@ -112,9 +112,8 @@ func (ts *IntegrationTestSuite) TearDownSuite() {
 			if last != nil {
 				ts.NoError(last)
 				return
-			} else {
-				ts.FailNow("leaks timed out but no error, should be impossible")
 			}
+			ts.FailNow("leaks timed out but no error, should be impossible")
 		case <-time.After(time.Second):
 			// https://github.com/uber-go/cadence-client/issues/739
 			last = goleak.FindLeaks(goleak.IgnoreTopFunction("go.uber.org/cadence/internal.(*coroutineState).initialYield"))