123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429 |
- 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, "/")
- }
|