package zookeeper import ( "fmt" "path" "strings" "time" log "github.com/Sirupsen/logrus" "github.com/docker/swarm/discovery" "github.com/samuel/go-zookeeper/zk" ) // DiscoveryService is exported type DiscoveryService struct { conn *zk.Conn path []string heartbeat uint64 } func init() { discovery.Register("zk", &DiscoveryService{}) } func (s *DiscoveryService) fullpath() string { return "/" + strings.Join(s.path, "/") } func (s *DiscoveryService) createFullpath() error { for i := 1; i <= len(s.path); i++ { newpath := "/" + strings.Join(s.path[:i], "/") _, err := s.conn.Create(newpath, []byte{1}, 0, zk.WorldACL(zk.PermAll)) if err != nil { // It's OK if key already existed. Just skip. if err != zk.ErrNodeExists { return err } } } return nil } // Initialize is exported func (s *DiscoveryService) Initialize(uris string, heartbeat uint64) error { var ( // split here because uris can contain multiples ips // like `zk://192.168.0.1,192.168.0.2,192.168.0.3/path` parts = strings.SplitN(uris, "/", 2) ips = strings.Split(parts[0], ",") ) if len(parts) != 2 { return fmt.Errorf("invalid format %q, missing ", uris) } if strings.Contains(parts[1], "/") { s.path = strings.Split(parts[1], "/") } else { s.path = []string{parts[1]} } conn, _, err := zk.Connect(ips, time.Second) if err != nil { return err } s.conn = conn s.heartbeat = heartbeat err = s.createFullpath() if err != nil { return err } return nil } // Fetch is exported func (s *DiscoveryService) Fetch() ([]*discovery.Entry, error) { addrs, _, err := s.conn.Children(s.fullpath()) if err != nil { return nil, err } return discovery.CreateEntries(addrs) } // Watch is exported func (s *DiscoveryService) Watch(callback discovery.WatchCallback) { addrs, _, eventChan, err := s.conn.ChildrenW(s.fullpath()) if err != nil { log.WithField("name", "zk").Debug("Discovery watch aborted") return } entries, err := discovery.CreateEntries(addrs) if err == nil { callback(entries) } for e := range eventChan { if e.Type == zk.EventNodeChildrenChanged { log.WithField("name", "zk").Debug("Discovery watch triggered") entries, err := s.Fetch() if err == nil { callback(entries) } } } } // Register is exported func (s *DiscoveryService) Register(addr string) error { nodePath := path.Join(s.fullpath(), addr) // check existing for the parent path first exist, _, err := s.conn.Exists(s.fullpath()) if err != nil { return err } // if the parent path does not exist yet if exist == false { // create the parent first err = s.createFullpath() if err != nil { return err } } else { // if node path exists exist, _, err = s.conn.Exists(nodePath) if err != nil { return err } // delete it first if exist { err = s.conn.Delete(nodePath, -1) if err != nil { return err } } } // create the node path to store address information _, err = s.conn.Create(nodePath, []byte(addr), 0, zk.WorldACL(zk.PermAll)) return err }