Add a unit test for ha proxy server

Author: Chao Xu

cmd/agent/main.go

@@ -64,8 +64,21 @@ type GrpcProxyAgentOptions struct {
 	proxyServerHost string
 	proxyServerPort int
 
-	agentID      string
-	syncInterval int
+	agentID           string
+	syncInterval      int
+	probeInterval     int
+	reconnectInterval int
+}
+
+func (o *GrpcProxyAgentOptions) ClientSetConfig(dialOption grpc.DialOption) *agentclient.ClientSetConfig {
+	return &agentclient.ClientSetConfig{
+		Address:           fmt.Sprintf("%s:%d", o.proxyServerHost, o.proxyServerPort),
+		AgentID:           o.agentID,
+		SyncInterval:      time.Duration(o.syncInterval) * time.Second,
+		ProbeInterval:     time.Duration(o.probeInterval) * time.Second,
+		ReconnectInterval: time.Duration(o.reconnectInterval) * time.Second,
+		DialOption:        dialOption,
+	}
 }
 
 func (o *GrpcProxyAgentOptions) Flags() *pflag.FlagSet {
@@ -77,6 +90,8 @@ func (o *GrpcProxyAgentOptions) Flags() *pflag.FlagSet {
 	flags.IntVar(&o.proxyServerPort, "proxy-server-port", o.proxyServerPort, "The port the proxy server is listening on.")
 	flags.StringVar(&o.agentID, "agent-id", o.agentID, "The unique ID of this agent. Default to a generated uuid if not set.")
 	flags.IntVar(&o.syncInterval, "sync-interval", o.syncInterval, "The seconds by which the agent periodically checks that it has connections to all instances of the proxy server.")
+	flags.IntVar(&o.probeInterval, "probe-interval", o.probeInterval, "The seconds by which the agent periodically checks if its connections to the proxy server are ready.")
+	flags.IntVar(&o.reconnectInterval, "reconnect-interval", o.reconnectInterval, "The seconds by which the agent tries to reconnect.")
 	return flags
 }
 
@@ -88,6 +103,8 @@ func (o *GrpcProxyAgentOptions) Print() {
 	klog.Warningf("ProxyServerPort set to %d.\n", o.proxyServerPort)
 	klog.Warningf("AgentID set to %s.\n", o.agentID)
 	klog.Warningf("SyncInterval set to %d.\n", o.syncInterval)
+	klog.Warningf("ProbeInterval set to %d.\n", o.probeInterval)
+	klog.Warningf("ReconnectInterval set to %d.\n", o.reconnectInterval)
 }
 
 func (o *GrpcProxyAgentOptions) Validate() error {
@@ -120,13 +137,15 @@ func (o *GrpcProxyAgentOptions) Validate() error {
 
 func newGrpcProxyAgentOptions() *GrpcProxyAgentOptions {
 	o := GrpcProxyAgentOptions{
-		agentCert:       "",
-		agentKey:        "",
-		caCert:          "",
-		proxyServerHost: "127.0.0.1",
-		proxyServerPort: 8091,
-		agentID:         uuid.New().String(),
-		syncInterval:    5,
+		agentCert:         "",
+		agentKey:          "",
+		caCert:            "",
+		proxyServerHost:   "127.0.0.1",
+		proxyServerPort:   8091,
+		agentID:           uuid.New().String(),
+		syncInterval:      5,
+		probeInterval:     5,
+		reconnectInterval: 5,
 	}
 	return &o
 }
@@ -173,7 +192,8 @@ func (a *Agent) runProxyConnection(o *GrpcProxyAgentOptions) error {
 		return err
 	}
 	dialOption := grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig))
-	cs := agentclient.NewAgentClientSet(fmt.Sprintf("%s:%d", o.proxyServerHost, o.proxyServerPort), o.agentID, time.Duration(o.syncInterval)*time.Second, dialOption)
+	cc := o.ClientSetConfig(dialOption)
+	cs := cc.NewAgentClientSet()
 	cs.Serve()
 
 	return nil

pkg/agent/agentclient/clientset.go

@@ -23,6 +23,7 @@ import (
 	"time"
 
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
 	"k8s.io/klog"
 )
 
@@ -40,14 +41,30 @@ type ClientSet struct {
 	syncInterval time.Duration // The interval by which the agent
 	// periodically checks that it has connections to all instances of the
 	// proxy server.
+	probeInterval time.Duration // The interval by which the agent
+	// periodically checks if its connections to the proxy server is ready.
+	reconnectInterval time.Duration // The interval by which the agent
+	// tries to reconnect.
 	dialOption grpc.DialOption
 }
 
-func (cs *ClientSet) clientsCount() int {
+func (cs *ClientSet) ClientsCount() int {
 	cs.mu.Lock()
 	defer cs.mu.Unlock()
 	return len(cs.clients)
 }
+func (cs *ClientSet) HealthyClientsCount() int {
+	cs.mu.Lock()
+	defer cs.mu.Unlock()
+	var count int
+	for _, c := range cs.clients {
+		if c.stream.conn.GetState() == connectivity.Ready {
+			count++
+		}
+	}
+	return count
+
+}
 
 func (cs *ClientSet) hasIDLocked(serverID string) bool {
 	_, ok := cs.clients[serverID]
@@ -81,14 +98,26 @@ func (cs *ClientSet) RemoveClient(serverID string) {
 	delete(cs.clients, serverID)
 }
 
-func NewAgentClientSet(address, agentID string, syncInterval time.Duration, dialOption grpc.DialOption) *ClientSet {
+type ClientSetConfig struct {
+	Address           string
+	AgentID           string
+	SyncInterval      time.Duration
+	ProbeInterval     time.Duration
+	ReconnectInterval time.Duration
+	DialOption        grpc.DialOption
+}
+
+func (cc *ClientSetConfig) NewAgentClientSet() *ClientSet {
 	return &ClientSet{
-		clients:      make(map[string]*AgentClient),
-		agentID:      agentID,
-		address:      address,
-		syncInterval: syncInterval,
-		dialOption:   dialOption,
+		clients:           make(map[string]*AgentClient),
+		agentID:           cc.AgentID,
+		address:           cc.Address,
+		syncInterval:      cc.SyncInterval,
+		probeInterval:     cc.ProbeInterval,
+		reconnectInterval: cc.ReconnectInterval,
+		dialOption:        cc.DialOption,
 	}
+
 }
 
 func (cs *ClientSet) newAgentClient() (*AgentClient, error) {
@@ -103,7 +132,7 @@ func (cs *ClientSet) sync() {
 			sleep := cs.syncInterval + time.Duration(rand.Float64()*jitter*float64(cs.syncInterval))
 			time.Sleep(sleep)
 		}
-		if cs.serverCount == 0 || cs.clientsCount() < cs.serverCount {
+		if cs.serverCount == 0 || cs.ClientsCount() < cs.serverCount {
 			c, err := cs.newAgentClient()
 			if err != nil {
 				klog.Error(err)

pkg/agent/agentclient/stream.go

@@ -33,7 +33,7 @@ import (
 )
 
 const (
-	defaultInterval = 5 * time.Second
+	defaultReconnectInterval = 5 * time.Second
 )
 
 type ReconnectError struct {
@@ -71,8 +71,8 @@ type RedialableAgentClient struct {
 	recvLock   sync.Mutex
 	reconnLock sync.Mutex
 
-	// Interval between every reconnect
-	Interval time.Duration
+	reconnectInterval time.Duration // interval between recoonects
+	probeInterval     time.Duration // interval between probe pings
 }
 
 func copyRedialableAgentClient(in RedialableAgentClient) RedialableAgentClient {
@@ -88,12 +88,13 @@ func copyRedialableAgentClient(in RedialableAgentClient) RedialableAgentClient {
 
 func NewRedialableAgentClient(address, agentID string, cs *ClientSet, opts ...grpc.DialOption) (*RedialableAgentClient, error) {
 	c := &RedialableAgentClient{
-		cs:       cs,
-		address:  address,
-		agentID:  agentID,
-		opts:     opts,
-		Interval: defaultInterval,
-		stopCh:   make(chan struct{}),
+		cs:                cs,
+		address:           address,
+		agentID:           agentID,
+		opts:              opts,
+		probeInterval:     cs.probeInterval,
+		reconnectInterval: cs.reconnectInterval,
+		stopCh:            make(chan struct{}),
 	}
 	serverID, err := c.Connect()
 	if err != nil {
@@ -108,7 +109,7 @@ func (c *RedialableAgentClient) probe() {
 		select {
 		case <-c.stopCh:
 			return
-		case <-time.After(c.Interval):
+		case <-time.After(c.probeInterval):
 			if c.conn == nil {
 				continue
 			}
@@ -253,8 +254,8 @@ func (c *RedialableAgentClient) reconnect() {
 		r, err := c.tryConnect()
 		if err != nil {
 			retry++
-			klog.Infof("Failed to connect to proxy server, retry %d in %v: %v", retry, c.Interval, err)
-			time.Sleep(c.Interval)
+			klog.Infof("Failed to connect to proxy server, retry %d in %v: %v", retry, c.reconnectInterval, err)
+			time.Sleep(c.reconnectInterval)
 			continue
 		}
 		switch {
@@ -272,8 +273,8 @@ func (c *RedialableAgentClient) reconnect() {
 				klog.Infof("failed to close connection to %s: %v", r.serverID, err)
 			}
 			retry++
-			klog.Infof("Trying to reconnect to proxy server %s, got connected to proxy server %s, for which there is already a connection, retry %d in %v", c.serverID, r.serverID, retry, c.Interval)
-			time.Sleep(c.Interval)
+			klog.Infof("Trying to reconnect to proxy server %s, got connected to proxy server %s, for which there is already a connection, retry %d in %v", c.serverID, r.serverID, retry, c.reconnectInterval)
+			time.Sleep(c.reconnectInterval)
 		case r.serverID != c.serverID && !c.cs.HasID(r.serverID):
 			// create a new client
 			cc := copyRedialableAgentClient(*c)
@@ -287,8 +288,8 @@ func (c *RedialableAgentClient) reconnect() {
 			}
 			go ac.Serve()
 			retry++
-			klog.Infof("Trying to reconnect to proxy server %s, got connected to proxy server %s. We will add this connection to the client set, but keep retrying connecting to proxy server %s, retry %d in %v", c.serverID, r.serverID, c.serverID, retry, c.Interval)
-			time.Sleep(c.Interval)
+			klog.Infof("Trying to reconnect to proxy server %s, got connected to proxy server %s. We will add this connection to the client set, but keep retrying connecting to proxy server %s, retry %d in %v", c.serverID, r.serverID, c.serverID, retry, c.reconnectInterval)
+			time.Sleep(c.reconnectInterval)
 		}
 	}
 

tests/concurrent_test.go

@@ -27,9 +27,7 @@ func TestProxy_Concurrency(t *testing.T) {
 	}
 	defer cleanup()
 
-	if err := runAgent(proxy.agent, stopCh); err != nil {
-		t.Fatal(err)
-	}
+	runAgent(proxy.agent, stopCh)
 
 	// Wait for agent to register on proxy server
 	time.Sleep(time.Second)