dnshostprovider.go

package zk

import (
	"fmt"
	"net"
	"sync"
)

// DNSHostProvider is the default HostProvider. It currently matches
// the Java StaticHostProvider, resolving hosts from DNS once during
// the call to Init.  It could be easily extended to re-query DNS
// periodically or if there is trouble connecting.
type DNSHostProvider struct {
	mu          sync.Mutex // Protects everything, so we can add asynchronous updates later.
	initServers []string   // init connect servers
	servers     []string
	curr        int
	last        int
	lookupHost  func(string) ([]string, error) // Override of net.LookupHost, for testing.
}

// Init is called first, with the servers specified in the connection
// string. It uses DNS to look up addresses for each server, then
// shuffles them all together.
func (hp *DNSHostProvider) Init(servers []string) error {
	hp.mu.Lock()
	defer hp.mu.Unlock()
	hp.initServers = make([]string, len(servers))
	copy(hp.initServers, servers)
	return hp.resolveServers(servers)
}

func (hp *DNSHostProvider) resolveServers(servers []string) error {
	lookupHost := hp.lookupHost
	if lookupHost == nil {
		lookupHost = net.LookupHost
	}

	found := []string{}
	for _, server := range servers {
		host, port, err := net.SplitHostPort(server)
		if err != nil {
			return err
		}
		addrs, err := lookupHost(host)
		if err != nil {
			return err
		}
		for _, addr := range addrs {
			found = append(found, net.JoinHostPort(addr, port))
		}
	}

	if len(found) == 0 {
		return fmt.Errorf("No hosts found for addresses %q", servers)
	}

	// Randomize the order of the servers to avoid creating hotspots
	stringShuffle(found)

	hp.servers = found
	hp.curr = -1
	hp.last = -1

	return nil
}

// Len returns the number of servers available
func (hp *DNSHostProvider) Len() int {
	hp.mu.Lock()
	defer hp.mu.Unlock()
	return len(hp.servers)
}

// Next returns the next server to connect to. retryStart will be true
// if we've looped through all known servers without Connected() being
// called.
func (hp *DNSHostProvider) Next() (server string, retryStart bool) {
	hp.mu.Lock()
	defer hp.mu.Unlock()
	hp.curr = (hp.curr + 1) % len(hp.servers)
	retryStart = hp.curr == hp.last
	if retryStart {
		if err := hp.resolveServers(hp.initServers); err != nil {
			return "", false
		}
		hp.curr = (hp.curr + 1) % len(hp.servers)
	}
	if hp.last == -1 {
		hp.last = 0
	}
	return hp.servers[hp.curr], false
}

// Connected notifies the HostProvider of a successful connection.
func (hp *DNSHostProvider) Connected() {
	hp.mu.Lock()
	defer hp.mu.Unlock()
	hp.last = hp.curr
}

func (hp *DNSHostProvider) GetServer() []string {
	hp.mu.Lock()
	defer hp.mu.Unlock()
	servers := make([]string, len(hp.servers))
	copy(servers, hp.servers)
	return servers
}