Fix data race (hashicorp#5029)

Fix hashicorp#4357
tinyspeck · Jan 8, 2019 · 238d430 · 238d430
1 parent 0670272
commit 238d430
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/watch/funcs.go b/watch/funcs.go
@@ -335,7 +335,7 @@ func agentServiceWatch(params map[string]interface{}) (WatcherFunc, error) {
 
 func makeQueryOptionsWithContext(p *Plan, stale bool) consulapi.QueryOptions {
 	ctx, cancel := context.WithCancel(context.Background())
-	p.cancelFunc = cancel
+	p.setCancelFunc(cancel)
 	opts := consulapi.QueryOptions{AllowStale: stale}
 	switch param := p.lastParamVal.(type) {
 	case WaitIndexVal:

diff --git a/watch/plan.go b/watch/plan.go
@@ -1,6 +1,7 @@
 package watch
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"os"
@@ -148,6 +149,17 @@ func (p *Plan) shouldStop() bool {
 	}
 }
 
+func (p *Plan) setCancelFunc(cancel context.CancelFunc) {
+	p.stopLock.Lock()
+	defer p.stopLock.Unlock()
+	if p.shouldStop() {
+		// The watch is stopped and execute the new cancel func to stop watchFactory
+		cancel()
+		return
+	}
+	p.cancelFunc = cancel
+}
+
 func (p *Plan) IsStopped() bool {
 	p.stopLock.Lock()
 	defer p.stopLock.Unlock()