Enhance HostProvider to update host provider on retryStart occurrence

conn.go

@@ -113,6 +113,8 @@ type Conn struct {
 	logInfo bool // true if information messages are logged; false if only errors are logged
 
 	buf []byte
+
+	configAddress []string
 }
 
 // connOption represents a connection option.
@@ -163,6 +165,9 @@ type HostProvider interface {
 	Next() (server string, retryStart bool)
 	// Notify the HostProvider of a successful connection.
 	Connected()
+	// UpdateServerList is called when HostProvider is abnormal.
+	// Update the list of servers.
+	UpdateServerList(servers []string) error
 }
 
 // ConnectWithDialer establishes a new connection to a pool of zookeeper servers
@@ -205,6 +210,7 @@ func Connect(servers []string, sessionTimeout time.Duration, options ...connOpti
 		logInfo:        true, // default is true for backwards compatability
 		buf:            make([]byte, bufferSize),
 		resendZkAuthFn: resendZkAuth,
+		configAddress:  servers,
 	}
 
 	// Set provided options.
@@ -365,6 +371,9 @@ func (c *Conn) sendEvent(evt Event) {
 
 func (c *Conn) connect() error {
 	var retryStart bool
+	// maybe need reload host provider when retryStart twice
+	var retryStartEd bool
+
 	for {
 		c.serverMu.Lock()
 		c.server, retryStart = c.hostProvider.Next()
@@ -373,14 +382,38 @@ func (c *Conn) connect() error {
 		c.setState(StateConnecting)
 
 		if retryStart {
-			c.flushUnsentRequests(ErrNoServer)
-			select {
-			case <-time.After(time.Second):
-				// pass
-			case <-c.shouldQuit:
-				c.setState(StateDisconnected)
-				c.flushUnsentRequests(ErrClosing)
-				return ErrClosing
+
+			if retryStartEd {
+				//hostProvider try to UpdateServerList
+				var err error
+				for {
+					err = c.hostProvider.UpdateServerList(c.configAddress)
+					if err != nil {
+						c.logger.Printf("host provider update server list err %v : %v", c.configAddress, err)
+						<-time.After(time.Second * 1)
+						continue
+					}
+					break
+				}
+
+				retryStartEd = false
+
+				c.serverMu.Lock()
+				c.server, retryStart = c.hostProvider.Next()
+				c.serverMu.Unlock()
+
+			} else {
+
+				c.flushUnsentRequests(ErrNoServer)
+				select {
+				case <-time.After(time.Second):
+					retryStartEd = true
+					// pass
+				case <-c.shouldQuit:
+					c.setState(StateDisconnected)
+					c.flushUnsentRequests(ErrClosing)
+					return ErrClosing
+				}
 			}
 		}
 

conn_test.go

@@ -2,8 +2,10 @@ package zk
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io/ioutil"
+	"net"
 	"sync"
 	"testing"
 	"time"
@@ -196,3 +198,69 @@ func TestNotifyWatches(t *testing.T) {
 		})
 	}
 }
+
+func TestConnLoop(t *testing.T) {
+
+	hostProvider := &DNSHostProvider{}
+
+	seq := 1
+	executedCount := 0
+	ipNum := 3
+
+	// simulate dns resolution to different ip addresses
+	hostProvider.lookupHost = func(string) ([]string, error) {
+		fmt.Println("Executed lookupHost,", executedCount)
+
+		executedCount++
+
+		if executedCount > 3 {
+			return nil, errors.New("lookupHost error")
+		}
+
+		ips := []string{}
+
+		for i := 0; i < ipNum; i++ {
+			seq = seq + 1
+			ips = append(ips, fmt.Sprintf("127.0.0.%v", seq))
+		}
+
+		return ips, nil
+	}
+
+	srvs := []string{"zk-test:2181"}
+
+	conn := &Conn{
+		dialer:         net.DialTimeout,
+		hostProvider:   hostProvider,
+		conn:           nil,
+		state:          StateDisconnected,
+		shouldQuit:     make(chan struct{}),
+		connectTimeout: 1 * time.Second,
+		sendChan:       make(chan *request, sendChanSize),
+		requests:       make(map[int32]*request),
+		watchers:       make(map[watchPathType][]chan Event),
+		passwd:         emptyPassword,
+		logger:         DefaultLogger,
+		logInfo:        true, // default is true for backwards compatability
+		buf:            make([]byte, bufferSize),
+		resendZkAuthFn: resendZkAuth,
+		configAddress:  srvs,
+	}
+
+	if err := conn.hostProvider.Init(srvs); err != nil {
+		t.Fatal(err)
+	}
+
+	for {
+		if err := conn.connect(); err != nil {
+			// c.Close() was called
+			return
+		}
+		conn.hostProvider.Connected()
+
+		// mock conn disconnect
+		<-time.After(time.Second * 3)
+
+		fmt.Println("conn disconnect")
+	}
+}

dnshostprovider.go

@@ -86,3 +86,9 @@ func (hp *DNSHostProvider) Connected() {
 	defer hp.mu.Unlock()
 	hp.last = hp.curr
 }
+
+// UpdateServerList is called when HostProvider is abnormal.
+// Update the list of servers.
+func (hp *DNSHostProvider) UpdateServerList(servers []string) error {
+	return hp.Init(servers)
+}

dnshostprovider_test.go

@@ -71,6 +71,9 @@ func newLocalHostPortsFacade(inner HostProvider, ports []int) *localHostPortsFac
 func (lhpf *localHostPortsFacade) Len() int                    { return lhpf.inner.Len() }
 func (lhpf *localHostPortsFacade) Connected()                  { lhpf.inner.Connected() }
 func (lhpf *localHostPortsFacade) Init(servers []string) error { return lhpf.inner.Init(servers) }
+func (lhpf *localHostPortsFacade) UpdateServerList(servers []string) error {
+	return lhpf.inner.UpdateServerList(servers)
+}
 func (lhpf *localHostPortsFacade) Next() (string, bool) {
 	server, retryStart := lhpf.inner.Next()