agentclient handles multiple waiters for reconnect

When reconnect happens, if there are multiple waiters, all of them
should get notification when the reconnect is done.

AgentClient differentiates between first connect with reconnect. If the
agent fails to connect to agent server, it is considered to be critical,
and we don't retry. After connection is established, we retry connecting
to server to recover from server failure.

Author: anfernee

cmd/agent/main.go

@@ -178,7 +178,10 @@ func (p *Agent) runProxyConnection(o *GrpcProxyAgentOptions) error {
 		RootCAs:      certPool,
 	})
 	dialOption := grpc.WithTransportCredentials(transportCreds)
-	client := agentclient.NewAgentClient(fmt.Sprintf("%s:%d", o.proxyServerHost, o.proxyServerPort), dialOption)
+	client, err := agentclient.NewAgentClient(fmt.Sprintf("%s:%d", o.proxyServerHost, o.proxyServerPort), dialOption)
+	if err != nil {
+		return err
+	}
 
 	stopCh := make(chan struct{})
 

pkg/agent/agentclient/client.go

@@ -37,13 +37,17 @@ type AgentClient struct {
 }
 
 // NewAgentClient creates an AgentClient
-func NewAgentClient(address string, opts ...grpc.DialOption) *AgentClient {
+func NewAgentClient(address string, opts ...grpc.DialOption) (*AgentClient, error) {
+	stream, err := NewRedialableAgentClient(address, opts...)
+	if err != nil {
+		return nil, err
+	}
+
 	a := &AgentClient{
 		connContext: make(map[int64]*connContext),
-		stream:      NewRedialableAgentClient(address, opts...),
+		stream:      stream,
 	}
-
-	return a
+	return a, nil
 }
 
 // connContext tracks a connection from agent to node network.

pkg/agent/agentclient/client_test.go

@@ -16,8 +16,10 @@ import (
 
 func TestServeData_HTTP(t *testing.T) {
 	var stream agent.AgentService_ConnectClient
-	client := NewAgentClient("")
-	client.stream, stream = pipe()
+	client := &AgentClient{
+		connContext: make(map[int64]*connContext),
+	}
+	client.stream, stream = pipe2()
 	stopCh := make(chan struct{})
 
 	// Start agent
@@ -100,8 +102,10 @@ func TestServeData_HTTP(t *testing.T) {
 
 func TestClose_Client(t *testing.T) {
 	var stream agent.AgentService_ConnectClient
-	client := NewAgentClient("")
-	client.stream, stream = pipe()
+	client := &AgentClient{
+		connContext: make(map[int64]*connContext),
+	}
+	client.stream, stream = pipe2()
 	stopCh := make(chan struct{})
 
 	// Start agent
@@ -193,6 +197,11 @@ func pipe() (agent.AgentService_ConnectClient, agent.AgentService_ConnectClient)
 	return s1, s2
 }
 
+func pipe2() (*RedialableAgentClient, agent.AgentService_ConnectClient) {
+	s1, s2 := pipe()
+	return &RedialableAgentClient{stream: s1}, s2
+}
+
 func (s *fakeStream) Send(packet *agent.Packet) error {
 	klog.Infof("[DEBUG] send packet %+v", packet)
 	s.w <- packet

pkg/agent/agentclient/stream.go

@@ -40,7 +40,8 @@ type RedialableAgentClient struct {
 	opts          []grpc.DialOption
 	conn          *grpc.ClientConn
 	stopCh        chan struct{}
-	reconnTrigger chan error
+	reconnOngoing bool
+	reconnWaiters []chan error
 
 	// locks
 	sendLock   sync.Mutex
@@ -54,7 +55,7 @@ type RedialableAgentClient struct {
 	Interval time.Duration
 }
 
-func NewRedialableAgentClient(address string, opts ...grpc.DialOption) *RedialableAgentClient {
+func NewRedialableAgentClient(address string, opts ...grpc.DialOption) (*RedialableAgentClient, error) {
 	c := &RedialableAgentClient{
 		address:  address,
 		opts:     opts,
@@ -63,11 +64,7 @@ func NewRedialableAgentClient(address string, opts ...grpc.DialOption) *Redialab
 		stopCh:   make(chan struct{}),
 	}
 
-	_ = <-c.triggerReconnect()
-
-	go c.probe()
-
-	return c
+	return c, c.Connect()
 }
 
 func (c *RedialableAgentClient) probe() {
@@ -129,22 +126,26 @@ func (c *RedialableAgentClient) triggerReconnect() <-chan error {
 	c.reconnLock.Lock()
 	defer c.reconnLock.Unlock()
 
-	if c.reconnTrigger != nil {
-		return c.reconnTrigger
-	}
+	errch := make(chan error)
+	c.reconnWaiters = append(c.reconnWaiters, errch)
 
-	c.reconnTrigger = make(chan error)
-	go c.reconnect()
+	if !c.reconnOngoing {
+		go c.reconnect()
+		c.reconnOngoing = true
+	}
 
-	return c.reconnTrigger
+	return errch
 }
 
 func (c *RedialableAgentClient) doneReconnect(err error) {
 	c.reconnLock.Lock()
 	defer c.reconnLock.Unlock()
 
-	c.reconnTrigger <- err
-	c.reconnTrigger = nil
+	for _, ch := range c.reconnWaiters {
+		ch <- err
+	}
+	c.reconnOngoing = false
+	c.reconnWaiters = nil
 }
 
 func (c *RedialableAgentClient) Recv() (*agent.Packet, error) {
@@ -164,6 +165,15 @@ func (c *RedialableAgentClient) Recv() (*agent.Packet, error) {
 	}
 }
 
+func (c *RedialableAgentClient) Connect() error {
+	if err := c.tryConnect(); err != nil {
+		return err
+	}
+
+	go c.probe()
+	return nil
+}
+
 func (c *RedialableAgentClient) reconnect() {
 	klog.Info("start to connect...")
 

pkg/agent/agentclient/stream_test.go

@@ -0,0 +1,124 @@
+package agentclient
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc"
+	"k8s.io/klog"
+	"sigs.k8s.io/apiserver-network-proxy/proto/agent"
+)
+
+func TestReconnectExits(t *testing.T) {
+	server := newTestServer("localhost:8899") // random addr
+	server.Start()
+	defer server.Stop()
+
+	time.Sleep(time.Millisecond)
+
+	client, err := NewRedialableAgentClient("localhost:8899", grpc.WithInsecure())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = client.Send(&agent.Packet{
+		Type: agent.PacketType_DIAL_REQ,
+	})
+	if err != nil {
+		t.Error(err)
+	}
+
+	client1 := make(chan bool)
+	go func() {
+		_, err := client.Recv()
+		if err != nil {
+			if err2, ok := err.(*ReconnectError); ok {
+				err2.Wait()
+				client1 <- true
+			}
+		}
+	}()
+
+	client2 := make(chan bool)
+	go func() {
+		_, err := client.Recv()
+		if err != nil {
+			if err2, ok := err.(*ReconnectError); ok {
+				err2.Wait()
+				client2 <- true
+			}
+		}
+	}()
+
+	client.interrupt()
+
+	var got1 bool
+	var got2 bool
+	select {
+	case got1 = <-client1:
+	case <-time.After(time.Second):
+	}
+	select {
+	case got2 = <-client2:
+	case <-time.After(time.Second):
+	}
+
+	if !got1 || !got2 {
+		t.Errorf("expect both clients get unblocked; not they don't (%t %t)", got1, got2)
+	}
+}
+
+type testServer struct {
+	addr       string
+	grpcServer *grpc.Server
+}
+
+func newTestServer(addr string) *testServer {
+	return &testServer{addr: addr}
+}
+
+func (s *testServer) Connect(stream agent.AgentService_ConnectServer) error {
+	stopCh := make(chan error)
+
+	// Recv only
+	go func() {
+		for {
+			_, err := stream.Recv()
+			if err == io.EOF {
+				close(stopCh)
+				return
+			}
+			if err != nil {
+				klog.Warningf(">>> Stream read from frontend error: %v", err)
+				close(stopCh)
+				return
+			}
+		}
+	}()
+
+	return <-stopCh
+}
+
+func (s *testServer) Start() error {
+	s.grpcServer = grpc.NewServer()
+	agent.RegisterAgentServiceServer(s.grpcServer, s)
+	lis, err := net.Listen("tcp", s.addr)
+	if err != nil {
+		return fmt.Errorf("failed to listen on %s: %v", s.addr, err)
+	}
+	go s.grpcServer.Serve(lis)
+	return nil
+}
+
+func (s *testServer) Stop() {
+	if s.grpcServer != nil {
+		s.grpcServer.Stop()
+	}
+}
+
+func (s *testServer) Addr() string {
+	return s.addr
+}

tests/proxy_test.go

@@ -37,7 +37,9 @@ func TestBasicProxy_GRPC(t *testing.T) {
 	proxy, cleanup, err := runGRPCProxyServer()
 	defer cleanup()
 
-	runAgent(proxy.agent, stopCh)
+	if err := runAgent(proxy.agent, stopCh); err != nil {
+		t.Fatal(err)
+	}
 
 	// Wait for agent to register on proxy server
 	time.Sleep(time.Second)
@@ -79,7 +81,9 @@ func TestBasicProxy_HTTPCONN(t *testing.T) {
 	proxy, cleanup, err := runHTTPConnProxyServer()
 	defer cleanup()
 
-	runAgent(proxy.agent, stopCh)
+	if err := runAgent(proxy.agent, stopCh); err != nil {
+		t.Fatal(err)
+	}
 
 	// Wait for agent to register on proxy server
 	time.Sleep(time.Second)
@@ -226,8 +230,13 @@ func runHTTPConnProxyServer() (proxy, func(), error) {
 	return proxy, cleanup, nil
 }
 
-func runAgent(addr string, stopCh <-chan struct{}) {
-	client := agentclient.NewAgentClient(addr, grpc.WithInsecure())
+func runAgent(addr string, stopCh <-chan struct{}) error {
+	client, err := agentclient.NewAgentClient(addr, grpc.WithInsecure())
+	if err != nil {
+		return err
+	}
 
 	go client.Serve(stopCh)
+
+	return nil
 }