hostdiscovery.go

package hostdiscovery

import (
	"net"
	"sync"

	"github.com/sirupsen/logrus"

	mapset "github.com/deckarep/golang-set"
	"github.com/docker/docker/pkg/discovery"
	// Including KV
	_ "github.com/docker/docker/pkg/discovery/kv"
	"github.com/docker/libkv/store/consul"
	"github.com/docker/libkv/store/etcd"
	"github.com/docker/libkv/store/zookeeper"
	"github.com/docker/libnetwork/types"
)

type hostDiscovery struct {
	watcher  discovery.Watcher
	nodes    mapset.Set
	stopChan chan struct{}
	sync.Mutex
}

func init() {
	consul.Register()
	etcd.Register()
	zookeeper.Register()
}

// NewHostDiscovery function creates a host discovery object
func NewHostDiscovery(watcher discovery.Watcher) HostDiscovery {
	return &hostDiscovery{watcher: watcher, nodes: mapset.NewSet(), stopChan: make(chan struct{})}
}

func (h *hostDiscovery) Watch(activeCallback ActiveCallback, joinCallback JoinCallback, leaveCallback LeaveCallback) error {
	h.Lock()
	d := h.watcher
	h.Unlock()
	if d == nil {
		return types.BadRequestErrorf("invalid discovery watcher")
	}
	discoveryCh, errCh := d.Watch(h.stopChan)
	go h.monitorDiscovery(discoveryCh, errCh, activeCallback, joinCallback, leaveCallback)
	return nil
}

func (h *hostDiscovery) monitorDiscovery(ch <-chan discovery.Entries, errCh <-chan error,
	activeCallback ActiveCallback, joinCallback JoinCallback, leaveCallback LeaveCallback) {
	for {
		select {
		case entries := <-ch:
			h.processCallback(entries, activeCallback, joinCallback, leaveCallback)
		case err := <-errCh:
			if err != nil {
				logrus.Errorf("discovery error: %v", err)
			}
		case <-h.stopChan:
			return
		}
	}
}

func (h *hostDiscovery) StopDiscovery() error {
	h.Lock()
	stopChan := h.stopChan
	h.watcher = nil
	h.Unlock()

	close(stopChan)
	return nil
}

func (h *hostDiscovery) processCallback(entries discovery.Entries,
	activeCallback ActiveCallback, joinCallback JoinCallback, leaveCallback LeaveCallback) {
	updated := hosts(entries)
	h.Lock()
	existing := h.nodes
	added, removed := diff(existing, updated)
	h.nodes = updated
	h.Unlock()

	activeCallback()
	if len(added) > 0 {
		joinCallback(added)
	}
	if len(removed) > 0 {
		leaveCallback(removed)
	}
}

func diff(existing mapset.Set, updated mapset.Set) (added []net.IP, removed []net.IP) {
	addSlice := updated.Difference(existing).ToSlice()
	removeSlice := existing.Difference(updated).ToSlice()
	for _, ip := range addSlice {
		added = append(added, net.ParseIP(ip.(string)))
	}
	for _, ip := range removeSlice {
		removed = append(removed, net.ParseIP(ip.(string)))
	}
	return
}

func (h *hostDiscovery) Fetch() []net.IP {
	h.Lock()
	defer h.Unlock()
	ips := []net.IP{}
	for _, ipstr := range h.nodes.ToSlice() {
		ips = append(ips, net.ParseIP(ipstr.(string)))
	}
	return ips
}

func hosts(entries discovery.Entries) mapset.Set {
	hosts := mapset.NewSet()
	for _, entry := range entries {
		hosts.Add(entry.Host)
	}
	return hosts
}