From 238d4302759f9b4ec2b580367aca1e5198450c90 Mon Sep 17 00:00:00 2001
From: dawxy <97687341@qq.com>
Date: Wed, 9 Jan 2019 00:43:14 +0800
Subject: [PATCH] Fix data race (#5029)

Fix #4357
---
 watch/funcs.go |  2 +-
 watch/plan.go  | 12 ++++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/watch/funcs.go b/watch/funcs.go
index 497ed7eb6e..3ff5e53095 100644
--- 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
index 1e34e4eacf..f0e8a6832a 100644
--- 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()