github.com/jfrazelle/docker@v1.1.2-0.20210712172922-bf78e25fe508/libnetwork/hostdiscovery/hostdiscovery.go (about)

     1  package hostdiscovery
     2  
     3  import (
     4  	"net"
     5  	"sync"
     6  
     7  	"github.com/sirupsen/logrus"
     8  
     9  	mapset "github.com/deckarep/golang-set"
    10  	"github.com/docker/docker/pkg/discovery"
    11  
    12  	// Including KV
    13  	"github.com/docker/docker/libnetwork/types"
    14  	_ "github.com/docker/docker/pkg/discovery/kv" // register all the things with host discovery
    15  	"github.com/docker/libkv/store/consul"
    16  	"github.com/docker/libkv/store/etcd"
    17  	"github.com/docker/libkv/store/zookeeper"
    18  )
    19  
    20  type hostDiscovery struct {
    21  	watcher  discovery.Watcher
    22  	nodes    mapset.Set
    23  	stopChan chan struct{}
    24  	sync.Mutex
    25  }
    26  
    27  func init() {
    28  	consul.Register()
    29  	etcd.Register()
    30  	zookeeper.Register()
    31  }
    32  
    33  // NewHostDiscovery function creates a host discovery object
    34  func NewHostDiscovery(watcher discovery.Watcher) HostDiscovery {
    35  	return &hostDiscovery{watcher: watcher, nodes: mapset.NewSet(), stopChan: make(chan struct{})}
    36  }
    37  
    38  func (h *hostDiscovery) Watch(activeCallback ActiveCallback, joinCallback JoinCallback, leaveCallback LeaveCallback) error {
    39  	h.Lock()
    40  	d := h.watcher
    41  	h.Unlock()
    42  	if d == nil {
    43  		return types.BadRequestErrorf("invalid discovery watcher")
    44  	}
    45  	discoveryCh, errCh := d.Watch(h.stopChan)
    46  	go h.monitorDiscovery(discoveryCh, errCh, activeCallback, joinCallback, leaveCallback)
    47  	return nil
    48  }
    49  
    50  func (h *hostDiscovery) monitorDiscovery(ch <-chan discovery.Entries, errCh <-chan error,
    51  	activeCallback ActiveCallback, joinCallback JoinCallback, leaveCallback LeaveCallback) {
    52  	for {
    53  		select {
    54  		case entries := <-ch:
    55  			h.processCallback(entries, activeCallback, joinCallback, leaveCallback)
    56  		case err := <-errCh:
    57  			if err != nil {
    58  				logrus.Errorf("discovery error: %v", err)
    59  			}
    60  		case <-h.stopChan:
    61  			return
    62  		}
    63  	}
    64  }
    65  
    66  func (h *hostDiscovery) StopDiscovery() error {
    67  	h.Lock()
    68  	stopChan := h.stopChan
    69  	h.watcher = nil
    70  	h.Unlock()
    71  
    72  	close(stopChan)
    73  	return nil
    74  }
    75  
    76  func (h *hostDiscovery) processCallback(entries discovery.Entries,
    77  	activeCallback ActiveCallback, joinCallback JoinCallback, leaveCallback LeaveCallback) {
    78  	updated := hosts(entries)
    79  	h.Lock()
    80  	existing := h.nodes
    81  	added, removed := diff(existing, updated)
    82  	h.nodes = updated
    83  	h.Unlock()
    84  
    85  	activeCallback()
    86  	if len(added) > 0 {
    87  		joinCallback(added)
    88  	}
    89  	if len(removed) > 0 {
    90  		leaveCallback(removed)
    91  	}
    92  }
    93  
    94  func diff(existing mapset.Set, updated mapset.Set) (added []net.IP, removed []net.IP) {
    95  	addSlice := updated.Difference(existing).ToSlice()
    96  	removeSlice := existing.Difference(updated).ToSlice()
    97  	for _, ip := range addSlice {
    98  		added = append(added, net.ParseIP(ip.(string)))
    99  	}
   100  	for _, ip := range removeSlice {
   101  		removed = append(removed, net.ParseIP(ip.(string)))
   102  	}
   103  	return
   104  }
   105  
   106  func (h *hostDiscovery) Fetch() []net.IP {
   107  	h.Lock()
   108  	defer h.Unlock()
   109  	ips := []net.IP{}
   110  	for _, ipstr := range h.nodes.ToSlice() {
   111  		ips = append(ips, net.ParseIP(ipstr.(string)))
   112  	}
   113  	return ips
   114  }
   115  
   116  func hosts(entries discovery.Entries) mapset.Set {
   117  	hosts := mapset.NewSet()
   118  	for _, entry := range entries {
   119  		hosts.Add(entry.Host)
   120  	}
   121  	return hosts
   122  }