Merge pull request #520 from mortent/FixStuckGoroutine

Avoid goroutine being blocked forever in WaitTask

pkg/apply/taskrunner/task.go

@@ -4,6 +4,7 @@
 package taskrunner
 
 import (
+	"context"
 	"fmt"
 	"reflect"
 	"time"
@@ -101,12 +102,16 @@ func (w *WaitTask) Start(taskContext *TaskContext) {
 // the WaitTask struct. Once the timer expires, it will send
 // a message on the EventChannel provided in the taskContext.
 func (w *WaitTask) setTimer(taskContext *TaskContext) {
-	timer := time.NewTimer(w.Timeout)
+	ctx, cancel := context.WithTimeout(context.Background(), w.Timeout)
 	go func() {
 		// TODO(mortent): See if there is a better way to do this. This
 		// solution will cause the goroutine to hang forever if the
 		// Timeout is cancelled.
-		<-timer.C
+		<-ctx.Done()
+		// If the context was cancelled, we don't want to send any events.
+		if ctx.Err() == context.Canceled {
+			return
+		}
 		select {
 		// We only send the TimeoutError to the eventChannel if no one has gotten
 		// to the token first.
@@ -130,7 +135,7 @@ func (w *WaitTask) setTimer(taskContext *TaskContext) {
 		}
 	}()
 	w.cancelFunc = func() {
-		timer.Stop()
+		cancel()
 	}
 }
 

pkg/apply/taskrunner/task_test.go

@@ -60,8 +60,10 @@ func TestWaitTask_TimeoutCancelled(t *testing.T) {
 	timer := time.NewTimer(3 * time.Second)
 
 	select {
+	case res := <-taskContext.EventChannel():
+		t.Errorf("didn't expect error on eventChannel, but got %v", res)
 	case res := <-taskContext.TaskChannel():
-		t.Errorf("didn't expect timeout error, but got %v", res.Err)
+		t.Errorf("didn't expect event on taskChannel, but got %v", res)
 	case <-timer.C:
 		return
 	}