store/zookeeper/zookeeper.go

package zookeeper

import (
	"strings"
	"time"

	"github.com/docker/libkv"
	"github.com/docker/libkv/store"
	zk "github.com/samuel/go-zookeeper/zk"
)

const (
	// SOH control character
	SOH = "\x01"

	defaultTimeout = 10 * time.Second
)

// Zookeeper is the receiver type for
// the Store interface
type Zookeeper struct {
	timeout time.Duration
	client  *zk.Conn
}

type zookeeperLock struct {
	client *zk.Conn
	lock   *zk.Lock
	key    string
	value  []byte
}

// Register registers zookeeper to libkv
func Register() {
	libkv.AddStore(store.ZK, New)
}

// New creates a new Zookeeper client given a
// list of endpoints and an optional tls config
func New(endpoints []string, options *store.Config) (store.Store, error) {
	s := &Zookeeper{}
	s.timeout = defaultTimeout

	// Set options
	if options != nil {
		if options.ConnectionTimeout != 0 {
			s.setTimeout(options.ConnectionTimeout)
		}
	}

	// Connect to Zookeeper
	conn, _, err := zk.Connect(endpoints, s.timeout)
	if err != nil {
		return nil, err
	}
	s.client = conn

	return s, nil
}

// setTimeout sets the timeout for connecting to Zookeeper
func (s *Zookeeper) setTimeout(time time.Duration) {
	s.timeout = time
}

// Get the value at "key", returns the last modified index
// to use in conjunction to Atomic calls
func (s *Zookeeper) Get(key string) (pair *store.KVPair, err error) {
	resp, meta, err := s.client.Get(s.normalize(key))

	if err != nil {
		if err == zk.ErrNoNode {
			return nil, store.ErrKeyNotFound
		}
		return nil, err
	}

	// FIXME handle very rare cases where Get returns the
	// SOH control character instead of the actual value
	if string(resp) == SOH {
		return s.Get(store.Normalize(key))
	}

	pair = &store.KVPair{
		Key:       key,
		Value:     resp,
		LastIndex: uint64(meta.Version),
	}

	return pair, nil
}

// createFullPath creates the entire path for a directory
// that does not exist
func (s *Zookeeper) createFullPath(path []string, ephemeral bool) error {
	for i := 1; i <= len(path); i++ {
		newpath := "/" + strings.Join(path[:i], "/")
		if i == len(path) && ephemeral {
			_, err := s.client.Create(newpath, []byte{}, zk.FlagEphemeral, zk.WorldACL(zk.PermAll))
			return err
		}
		_, err := s.client.Create(newpath, []byte{}, 0, zk.WorldACL(zk.PermAll))
		if err != nil {
			// Skip if node already exists
			if err != zk.ErrNodeExists {
				return err
			}
		}
	}
	return nil
}

// Put a value at "key"
func (s *Zookeeper) Put(key string, value []byte, opts *store.WriteOptions) error {
	fkey := s.normalize(key)

	exists, err := s.Exists(key)
	if err != nil {
		return err
	}

	if !exists {
		if opts != nil && opts.TTL > 0 {
			s.createFullPath(store.SplitKey(strings.TrimSuffix(key, "/")), true)
		} else {
			s.createFullPath(store.SplitKey(strings.TrimSuffix(key, "/")), false)
		}
	}

	_, err = s.client.Set(fkey, value, -1)
	return err
}

// Delete a value at "key"
func (s *Zookeeper) Delete(key string) error {
	err := s.client.Delete(s.normalize(key), -1)
	if err == zk.ErrNoNode {
		return store.ErrKeyNotFound
	}
	return err
}

// Exists checks if the key exists inside the store
func (s *Zookeeper) Exists(key string) (bool, error) {
	exists, _, err := s.client.Exists(s.normalize(key))
	if err != nil {
		return false, err
	}
	return exists, nil
}

// Watch for changes on a "key"
// It returns a channel that will receive changes or pass
// on errors. Upon creation, the current value will first
// be sent to the channel. Providing a non-nil stopCh can
// be used to stop watching.
func (s *Zookeeper) Watch(key string, stopCh <-chan struct{}) (<-chan *store.KVPair, error) {
	// Get the key first
	pair, err := s.Get(key)
	if err != nil {
		return nil, err
	}

	// Catch zk notifications and fire changes into the channel.
	watchCh := make(chan *store.KVPair)
	go func() {
		defer close(watchCh)

		// Get returns the current value to the channel prior
		// to listening to any event that may occur on that key
		watchCh <- pair
		for {
			_, _, eventCh, err := s.client.GetW(s.normalize(key))
			if err != nil {
				return
			}
			select {
			case e := <-eventCh:
				if e.Type == zk.EventNodeDataChanged {
					if entry, err := s.Get(key); err == nil {
						watchCh <- entry
					}
				}
			case <-stopCh:
				// There is no way to stop GetW so just quit
				return
			}
		}
	}()

	return watchCh, nil
}

// WatchTree watches for changes on a "directory"
// It returns a channel that will receive changes or pass
// on errors. Upon creating a watch, the current childs values
// will be sent to the channel .Providing a non-nil stopCh can
// be used to stop watching.
func (s *Zookeeper) WatchTree(directory string, stopCh <-chan struct{}) (<-chan []*store.KVPair, error) {
	// List the childrens first
	entries, err := s.List(directory)
	if err != nil {
		return nil, err
	}

	// Catch zk notifications and fire changes into the channel.
	watchCh := make(chan []*store.KVPair)
	go func() {
		defer close(watchCh)

		// List returns the children values to the channel
		// prior to listening to any events that may occur
		// on those keys
		watchCh <- entries

		for {
			_, _, eventCh, err := s.client.ChildrenW(s.normalize(directory))
			if err != nil {
				return
			}
			select {
			case e := <-eventCh:
				if e.Type == zk.EventNodeChildrenChanged {
					if kv, err := s.List(directory); err == nil {
						watchCh <- kv
					}
				}
			case <-stopCh:
				// There is no way to stop GetW so just quit
				return
			}
		}
	}()

	return watchCh, nil
}

// List child nodes of a given directory
func (s *Zookeeper) List(directory string) ([]*store.KVPair, error) {
	keys, stat, err := s.client.Children(s.normalize(directory))
	if err != nil {
		if err == zk.ErrNoNode {
			return nil, store.ErrKeyNotFound
		}
		return nil, err
	}

	kv := []*store.KVPair{}

	// FIXME Costly Get request for each child key..
	for _, key := range keys {
		pair, err := s.Get(strings.TrimSuffix(directory, "/") + s.normalize(key))
		if err != nil {
			// If node is not found: List is out of date, retry
			if err == store.ErrKeyNotFound {
				return s.List(directory)
			}
			return nil, err
		}

		kv = append(kv, &store.KVPair{
			Key:       key,
			Value:     []byte(pair.Value),
			LastIndex: uint64(stat.Version),
		})
	}

	return kv, nil
}

// DeleteTree deletes a range of keys under a given directory
func (s *Zookeeper) DeleteTree(directory string) error {
	pairs, err := s.List(directory)
	if err != nil {
		return err
	}

	var reqs []interface{}

	for _, pair := range pairs {
		reqs = append(reqs, &zk.DeleteRequest{
			Path:    s.normalize(directory + "/" + pair.Key),
			Version: -1,
		})
	}

	_, err = s.client.Multi(reqs...)
	return err
}

// AtomicPut put a value at "key" if the key has not been
// modified in the meantime, throws an error if this is the case
func (s *Zookeeper) AtomicPut(key string, value []byte, previous *store.KVPair, _ *store.WriteOptions) (bool, *store.KVPair, error) {
	var lastIndex uint64

	if previous != nil {
		meta, err := s.client.Set(s.normalize(key), value, int32(previous.LastIndex))
		if err != nil {
			// Compare Failed
			if err == zk.ErrBadVersion {
				return false, nil, store.ErrKeyModified
			}
			return false, nil, err
		}
		lastIndex = uint64(meta.Version)
	} else {
		// Interpret previous == nil as create operation.
		_, err := s.client.Create(s.normalize(key), value, 0, zk.WorldACL(zk.PermAll))
		if err != nil {
			// Directory does not exist
			if err == zk.ErrNoNode {

				// Create the directory
				parts := store.SplitKey(strings.TrimSuffix(key, "/"))
				parts = parts[:len(parts)-1]
				if err = s.createFullPath(parts, false); err != nil {
					// Failed to create the directory.
					return false, nil, err
				}

				// Create the node
				if _, err := s.client.Create(s.normalize(key), value, 0, zk.WorldACL(zk.PermAll)); err != nil {
					// Node exist error (when previous nil)
					if err == zk.ErrNodeExists {
						return false, nil, store.ErrKeyExists
					}
					return false, nil, err
				}

			} else {
				// Node Exists error (when previous nil)
				if err == zk.ErrNodeExists {
					return false, nil, store.ErrKeyExists
				}

				// Unhandled error
				return false, nil, err
			}
		}
		lastIndex = 0 // Newly created nodes have version 0.
	}

	pair := &store.KVPair{
		Key:       key,
		Value:     value,
		LastIndex: lastIndex,
	}

	return true, pair, nil
}

// AtomicDelete deletes a value at "key" if the key
// has not been modified in the meantime, throws an
// error if this is the case
func (s *Zookeeper) AtomicDelete(key string, previous *store.KVPair) (bool, error) {
	if previous == nil {
		return false, store.ErrPreviousNotSpecified
	}

	err := s.client.Delete(s.normalize(key), int32(previous.LastIndex))
	if err != nil {
		// Key not found
		if err == zk.ErrNoNode {
			return false, store.ErrKeyNotFound
		}
		// Compare failed
		if err == zk.ErrBadVersion {
			return false, store.ErrKeyModified
		}
		// General store error
		return false, err
	}
	return true, nil
}

// NewLock returns a handle to a lock struct which can
// be used to provide mutual exclusion on a key
func (s *Zookeeper) NewLock(key string, options *store.LockOptions) (lock store.Locker, err error) {
	value := []byte("")

	// Apply options
	if options != nil {
		if options.Value != nil {
			value = options.Value
		}
	}

	lock = &zookeeperLock{
		client: s.client,
		key:    s.normalize(key),
		value:  value,
		lock:   zk.NewLock(s.client, s.normalize(key), zk.WorldACL(zk.PermAll)),
	}

	return lock, err
}

// Lock attempts to acquire the lock and blocks while
// doing so. It returns a channel that is closed if our
// lock is lost or if an error occurs
func (l *zookeeperLock) Lock(stopChan chan struct{}) (<-chan struct{}, error) {
	err := l.lock.Lock()

	if err == nil {
		// We hold the lock, we can set our value
		// FIXME: The value is left behind
		// (problematic for leader election)
		_, err = l.client.Set(l.key, l.value, -1)
	}

	return make(chan struct{}), err
}

// Unlock the "key". Calling unlock while
// not holding the lock will throw an error
func (l *zookeeperLock) Unlock() error {
	return l.lock.Unlock()
}

// Close closes the client connection
func (s *Zookeeper) Close() {
	s.client.Close()
}

// Normalize the key for usage in Zookeeper
func (s *Zookeeper) normalize(key string) string {
	key = store.Normalize(key)
	return strings.TrimSuffix(key, "/")
}