From d04c177a107dc41f6d16d598a1457280c3955551 Mon Sep 17 00:00:00 2001
From: Chun Chen <ramichen@tencent.com>
Date: Mon, 17 Aug 2015 16:07:43 +0800
Subject: [PATCH] Fix for zookeeper backend

Signed-off-by: Chun Chen <ramichen@tencent.com>
---
 libnetwork/Makefile      |  9 ++++++---
 libnetwork/store.go      | 23 +++++++++++++++++++++--
 libnetwork/store_test.go | 20 ++++++++++++++++++++
 3 files changed, 47 insertions(+), 5 deletions(-)
 create mode 100644 libnetwork/store_test.go

diff --git a/libnetwork/Makefile b/libnetwork/Makefile
index fa5a4fe1ab..1cbccbf3c0 100644
--- a/libnetwork/Makefile
+++ b/libnetwork/Makefile
@@ -1,4 +1,4 @@
-.PHONY: all all-local build build-local check check-code check-format run-tests check-local integration-tests install-deps coveralls circle-ci
+.PHONY: all all-local build build-local check check-code check-format run-tests check-local integration-tests install-deps coveralls circle-ci start-services
 SHELL=/bin/bash
 build_image=libnetwork-build
 dockerargs = --privileged -v $(shell pwd):/go/src/github.com/docker/libnetwork -w /go/src/github.com/docker/libnetwork
@@ -67,10 +67,10 @@ run-tests:
 	done
 	@echo "Done running tests"
 
-check-local:	check-format check-code run-tests
+check-local: 	check-format check-code start-services run-tests
 
 install-deps:
-	apt-get update && apt-get -y install iptables
+	apt-get update && apt-get -y install iptables zookeeperd
 	git clone https://github.com/golang/tools /go/src/golang.org/x/tools
 	go install golang.org/x/tools/cmd/vet
 	go install golang.org/x/tools/cmd/goimports
@@ -88,3 +88,6 @@ coveralls:
 circle-ci:
 	@${cidocker} make install-deps build-local check-local coveralls
 	make integration-tests
+
+start-services:
+	service zookeeper start
diff --git a/libnetwork/store.go b/libnetwork/store.go
index e14aa72b01..42a14de332 100644
--- a/libnetwork/store.go
+++ b/libnetwork/store.go
@@ -174,7 +174,11 @@ func (c *controller) watchNetworks() error {
 	cs := c.store
 	c.Unlock()
 
-	nwPairs, err := cs.KVStore().WatchTree(datastore.Key(datastore.NetworkKeyPrefix), nil)
+	networkKey := datastore.Key(datastore.NetworkKeyPrefix)
+	if err := ensureKeys(networkKey, cs); err != nil {
+		return fmt.Errorf("failed to ensure if the network keys are valid and present in store: %v", err)
+	}
+	nwPairs, err := cs.KVStore().WatchTree(networkKey, nil)
 	if err != nil {
 		return err
 	}
@@ -228,7 +232,11 @@ func (n *network) watchEndpoints() error {
 	stopCh := n.stopWatchCh
 	n.Unlock()
 
-	epPairs, err := cs.KVStore().WatchTree(datastore.Key(tmp.KeyPrefix()...), stopCh)
+	endpointKey := datastore.Key(tmp.KeyPrefix()...)
+	if err := ensureKeys(endpointKey, cs); err != nil {
+		return fmt.Errorf("failed to ensure if the endpoint keys are valid and present in store: %v", err)
+	}
+	epPairs, err := cs.KVStore().WatchTree(endpointKey, stopCh)
 	if err != nil {
 		return err
 	}
@@ -362,3 +370,14 @@ func (c *controller) processEndpointUpdate(ep *endpoint) bool {
 
 	return false
 }
+
+func ensureKeys(key string, cs datastore.DataStore) error {
+	exists, err := cs.KVStore().Exists(key)
+	if err != nil {
+		return err
+	}
+	if exists {
+		return nil
+	}
+	return cs.KVStore().Put(key, []byte{}, nil)
+}
diff --git a/libnetwork/store_test.go b/libnetwork/store_test.go
new file mode 100644
index 0000000000..d97601ba4e
--- /dev/null
+++ b/libnetwork/store_test.go
@@ -0,0 +1,20 @@
+package libnetwork
+
+import (
+	"testing"
+
+	"github.com/docker/libnetwork/config"
+)
+
+func TestZooKeeperBackend(t *testing.T) {
+	testNewController(t, "zk", "127.0.0.1:2181")
+}
+
+func testNewController(t *testing.T, provider, url string) error {
+	netOptions := []config.Option{}
+	netOptions = append(netOptions, config.OptionKVProvider(provider))
+	netOptions = append(netOptions, config.OptionKVProviderURL(url))
+
+	_, err := New(netOptions...)
+	return err
+}