From 2c8ecd53bcc5c681d683b8f03c70e8df7af66053 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Wed, 6 May 2015 11:08:08 -0700
Subject: [PATCH] physical/zk: Style changes and more error checking

---
 physical/zookeeper.go | 65 ++++++++++++++++++++++++-------------------
 1 file changed, 37 insertions(+), 28 deletions(-)

diff --git a/physical/zookeeper.go b/physical/zookeeper.go
index 1dbb71a320..1b01f4cbc6 100644
--- a/physical/zookeeper.go
+++ b/physical/zookeeper.go
@@ -35,14 +35,14 @@ func newZookeeperBackend(conf map[string]string) (Backend, error) {
 		path = "/" + path
 	}
 
-	// Configure the client
+	// Configure the client, default to localhost instance
 	var machines string
 	machines, ok = conf["address"]
 	if !ok {
-		// Default to the localhost instance
 		machines = "localhost:2181"
 	}
 
+	// Attempt to create the ZK client
 	client, _, err := zk.Connect(strings.Split(machines, ","), time.Second)
 	if err != nil {
 		return nil, fmt.Errorf("client setup failed: %v", err)
@@ -56,62 +56,70 @@ func newZookeeperBackend(conf map[string]string) (Backend, error) {
 	return c, nil
 }
 
-// zookeeper requires nodes to be there before set and get
-func (c *ZookeeperBackend) ensurePath(path string, value []byte) {
-
+// ensurePath is used to create each node in the path hierarchy.
+// We avoid calling this optimistically, and invoke it when we get
+// an error during an operation
+func (c *ZookeeperBackend) ensurePath(path string, value []byte) error {
 	nodes := strings.Split(path, "/")
-
 	acl := zk.WorldACL(zk.PermAll)
-
 	fullPath := ""
 	for index, node := range nodes {
 		if strings.TrimSpace(node) != "" {
 			fullPath += "/" + node
-
 			isLastNode := index+1 == len(nodes)
 
 			// set parent nodes to nil, leaf to value
 			// this block reduces round trips by being smart on the leaf create/set
 			if exists, _, _ := c.client.Exists(fullPath); !isLastNode && !exists {
-				c.client.Create(fullPath, nil, int32(0), acl)
+				if _, err := c.client.Create(fullPath, nil, int32(0), acl); err != nil {
+					return err
+				}
 			} else if isLastNode && !exists {
-				c.client.Create(fullPath, value, int32(0), acl)
+				if _, err := c.client.Create(fullPath, value, int32(0), acl); err != nil {
+					return err
+				}
 			} else if isLastNode && exists {
-				c.client.Set(fullPath, value, int32(0))
+				if _, err := c.client.Set(fullPath, value, int32(0)); err != nil {
+					return err
+				}
 			}
-
 		}
 	}
+	return nil
 }
 
 // Put is used to insert or update an entry
 func (c *ZookeeperBackend) Put(entry *Entry) error {
 	defer metrics.MeasureSince([]string{"zookeeper", "put"}, time.Now())
 
+	// Attempt to set the full path
 	fullPath := c.path + entry.Key
-
 	_, err := c.client.Set(fullPath, entry.Value, 0)
 
+	// If we get ErrNoNode, we need to construct the path hierarchy
 	if err == zk.ErrNoNode {
-		c.ensurePath(fullPath, entry.Value)
-		return nil
-	} else {
-		return err
+		return c.ensurePath(fullPath, entry.Value)
 	}
+	return err
 }
 
 // Get is used to fetch an entry
 func (c *ZookeeperBackend) Get(key string) (*Entry, error) {
 	defer metrics.MeasureSince([]string{"zookeeper", "get"}, time.Now())
 
+	// Attempt to read the full path
 	fullPath := c.path + key
-
 	value, _, err := c.client.Get(fullPath)
 
-	// dont error if the node doesnt exist
-	if err != nil && err != zk.ErrNoNode {
+	// Ignore if the node does not exist
+	if err == zk.ErrNoNode {
+		err = nil
+	}
+	if err != nil {
 		return nil, err
 	}
+
+	// Handle a non-existing value
 	if value == nil {
 		return nil, nil
 	}
@@ -126,15 +134,15 @@ func (c *ZookeeperBackend) Get(key string) (*Entry, error) {
 func (c *ZookeeperBackend) Delete(key string) error {
 	defer metrics.MeasureSince([]string{"zookeeper", "delete"}, time.Now())
 
+	// Delete the full path
 	fullPath := c.path + key
-
 	err := c.client.Delete(fullPath, -1)
 
+	// Mask if the node does not exist
 	if err == zk.ErrNoNode {
-		return nil
-	} else {
-		return err
+		err = nil
 	}
+	return err
 }
 
 // List is used ot list all the keys under a given
@@ -142,26 +150,27 @@ func (c *ZookeeperBackend) Delete(key string) error {
 func (c *ZookeeperBackend) List(prefix string) ([]string, error) {
 	defer metrics.MeasureSince([]string{"zookeeper", "list"}, time.Now())
 
+	// Query the children at the full path
 	fullPath := strings.TrimSuffix(c.path+prefix, "/")
-
 	result, _, err := c.client.Children(fullPath)
 
+	// If the path nodes are missing, no children!
 	if err == zk.ErrNoNode {
 		return []string{}, nil
 	}
 
 	children := []string{}
-
 	for _, key := range result {
 		children = append(children, key)
 
+		// Check if this entry has any child entries,
+		// and append the slash which is what Vault depends on
+		// for iteration
 		nodeChildren, _, _ := c.client.Children(fullPath + "/" + key)
 		if nodeChildren != nil && len(nodeChildren) > 0 {
 			children = append(children, key+"/")
 		}
 	}
-
 	sort.Strings(children)
-
 	return children, nil
 }