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