From 60cf4015a4a6e86e2f990d41b7d40faff4709386 Mon Sep 17 00:00:00 2001
From: Julius Volz <julius@soundcloud.com>
Date: Tue, 6 Oct 2015 15:47:34 +0200
Subject: [PATCH] Fix watching multiple Zookeeper paths in serverset SD.

Fix https://github.com/prometheus/prometheus/issues/1137
---
 retrieval/discovery/serverset.go | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/retrieval/discovery/serverset.go b/retrieval/discovery/serverset.go
index dedc08730..0e1265616 100644
--- a/retrieval/discovery/serverset.go
+++ b/retrieval/discovery/serverset.go
@@ -63,13 +63,13 @@ func (zl zookeeperLogger) Printf(s string, i ...interface{}) {
 // ServersetDiscovery retrieves target information from a Serverset server
 // and updates them via watches.
 type ServersetDiscovery struct {
-	conf      *config.ServersetSDConfig
-	conn      *zk.Conn
-	mu        sync.RWMutex
-	sources   map[string]*config.TargetGroup
-	sdUpdates *chan<- *config.TargetGroup
-	updates   chan zookeeperTreeCacheEvent
-	treeCache *zookeeperTreeCache
+	conf       *config.ServersetSDConfig
+	conn       *zk.Conn
+	mu         sync.RWMutex
+	sources    map[string]*config.TargetGroup
+	sdUpdates  *chan<- *config.TargetGroup
+	updates    chan zookeeperTreeCacheEvent
+	treeCaches []*zookeeperTreeCache
 }
 
 // NewServersetDiscovery returns a new ServersetDiscovery for the given config.
@@ -87,7 +87,9 @@ func NewServersetDiscovery(conf *config.ServersetSDConfig) *ServersetDiscovery {
 		sources: map[string]*config.TargetGroup{},
 	}
 	go sd.processUpdates()
-	sd.treeCache = newZookeeperTreeCache(conn, conf.Paths[0], updates)
+	for _, path := range conf.Paths {
+		sd.treeCaches = append(sd.treeCaches, newZookeeperTreeCache(conn, path, updates))
+	}
 	return sd
 }
 
@@ -143,7 +145,9 @@ func (sd *ServersetDiscovery) Run(ch chan<- *config.TargetGroup, done <-chan str
 	sd.mu.Unlock()
 
 	<-done
-	sd.treeCache.Stop()
+	for _, tc := range sd.treeCaches {
+		tc.Stop()
+	}
 }
 
 func parseServersetMember(data []byte, path string) (*model.LabelSet, error) {