From 0b9fca983b05336e5996c08afecbcfcc7324fdfe Mon Sep 17 00:00:00 2001
From: Stephan Erb <github@stephanerb.eu>
Date: Wed, 3 May 2017 01:21:37 +0200
Subject: [PATCH] Fix reload of ZooKeeper service discovery config (#2669)

Rational:

* When the config is reloaded and the provider context is canceled, we need to
  exit the current ZK `TargetProvider.Run` method as a new provider will be
  instantiated.
* In case `Stop` is called on the `ZookeeperTreeCache`, the update/events
  channel may not be closed as it is shared by multiple caches and would
  thus be double closed.
* Stopping all `zookeeperTreeCacheNode`s on teardown ensures all associated
  watcher go-routines will be closed eagerly rather than implicityly on
  connection close events.
---
 discovery/zookeeper/zookeeper.go |  1 +
 util/treecache/treecache.go      | 12 +++++++++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/discovery/zookeeper/zookeeper.go b/discovery/zookeeper/zookeeper.go
index 205eff726..0b8f1943b 100644
--- a/discovery/zookeeper/zookeeper.go
+++ b/discovery/zookeeper/zookeeper.go
@@ -93,6 +93,7 @@ func (d *Discovery) Run(ctx context.Context, ch chan<- []*config.TargetGroup) {
 	for {
 		select {
 		case <-ctx.Done():
+			return
 		case event := <-d.updates:
 			tg := &config.TargetGroup{
 				Source: event.Path,
diff --git a/util/treecache/treecache.go b/util/treecache/treecache.go
index 6f0e213b1..43af4cf6d 100644
--- a/util/treecache/treecache.go
+++ b/util/treecache/treecache.go
@@ -168,7 +168,7 @@ func (tc *ZookeeperTreeCache) loop(path string) {
 				failureMode = false
 			}
 		case <-tc.stop:
-			close(tc.events)
+			tc.recursiveStop(tc.head)
 			return
 		}
 	}
@@ -264,3 +264,13 @@ func (tc *ZookeeperTreeCache) recursiveDelete(path string, node *zookeeperTreeCa
 		tc.recursiveDelete(path+"/"+name, childNode)
 	}
 }
+
+func (tc *ZookeeperTreeCache) recursiveStop(node *zookeeperTreeCacheNode) {
+	if !node.stopped {
+		node.done <- struct{}{}
+		node.stopped = true
+	}
+	for _, childNode := range node.children {
+		tc.recursiveStop(childNode)
+	}
+}