Merge pull request #31 from anfernee/agent

Agent supports reconnect

Author: cheftako

.gitignore

@@ -1 +1,7 @@
 /bin
+/certs/
+/cfssl
+/cfssljson
+/easy-rsa-master/
+/easy-rsa.tar.gz
+

Makefile

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-.PHONY: gen clean certs build docker/proxy-server docker/proxy-agent push-images
+.PHONY: gen clean certs build docker/proxy-server docker/proxy-agent push-images test
 proto/agent/agent.pb.go: proto/agent/agent.proto
 	protoc -I proto proto/agent/agent.proto --go_out=plugins=grpc:proto
 
@@ -104,3 +104,6 @@ push-images: docker/proxy-agent docker/proxy-server
 
 clean:
 	rm -rf proto/agent/agent.pb.go proto/proxy.pb.go easy-rsa.tar.gz easy-rsa-master cfssl cfssljson certs bin
+
+test:
+	go test ./...

cmd/agent/main.go

@@ -30,9 +30,9 @@ import (
 	"github.com/spf13/pflag"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
+	"k8s.io/klog"
 	"sigs.k8s.io/apiserver-network-proxy/pkg/agent/agentclient"
 	"sigs.k8s.io/apiserver-network-proxy/pkg/util"
-	"k8s.io/klog"
 )
 
 func main() {
@@ -113,9 +113,9 @@ func (o *GrpcProxyAgentOptions) Validate() error {
 
 func newGrpcProxyAgentOptions() *GrpcProxyAgentOptions {
 	o := GrpcProxyAgentOptions{
-		agentCert: "",
-		agentKey:  "",
-		caCert:    "",
+		agentCert:       "",
+		agentKey:        "",
+		caCert:          "",
 		proxyServerHost: "127.0.0.1",
 		proxyServerPort: 8091,
 	}
@@ -178,10 +178,9 @@ func (p *Agent) runProxyConnection(o *GrpcProxyAgentOptions) error {
 		RootCAs:      certPool,
 	})
 	dialOption := grpc.WithTransportCredentials(transportCreds)
-	client := agentclient.NewAgentClient(fmt.Sprintf("%s:%d", o.proxyServerHost, o.proxyServerPort))
-
-	if err := client.Connect(dialOption); err != nil {
-		return fmt.Errorf("failed to connect to proxy-server: %v", err)
+	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

@@ -17,7 +17,6 @@ limitations under the License.
 package agentclient
 
 import (
-	"context"
 	"io"
 	"net"
 	"sync"
@@ -33,19 +32,22 @@ import (
 type AgentClient struct {
 	nextConnID  int64
 	connContext map[int64]*connContext
-	address     string
 
-	stream agent.AgentService_ConnectClient
+	stream *RedialableAgentClient
 }
 
 // NewAgentClient creates an AgentClient
-func NewAgentClient(address string) *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),
-		address:     address,
+		stream:      stream,
 	}
-
-	return a
+	return a, nil
 }
 
 // connContext tracks a connection from agent to node network.
@@ -68,21 +70,6 @@ func (c *connContext) cleanup() {
 //
 // The caller needs to call Serve to start serving proxy requests
 // coming from proxy server.
-func (a *AgentClient) Connect(opts ...grpc.DialOption) error {
-	c, err := grpc.Dial(a.address, opts...)
-	if err != nil {
-		return err
-	}
-
-	client := agent.NewAgentServiceClient(c)
-
-	a.stream, err = client.Connect(context.Background())
-	if err != nil {
-		return err
-	}
-
-	return nil
-}
 
 // Serve starts to serve proxied requests from proxy server over the
 // gRPC stream. Successful Connect is required before Serve. The
@@ -98,10 +85,16 @@ func (a *AgentClient) Serve(stopCh <-chan struct{}) {
 		}
 
 		pkt, err := a.stream.Recv()
-		if err == io.EOF {
-			klog.Info("received EOF, exit")
-			return
+		if err != nil {
+			if err2, ok := err.(*ReconnectError); ok {
+				err = err2.Wait()
+				continue
+			} else if err == io.EOF {
+				klog.Info("received EOF, exit")
+				return
+			}
 		}
+
 		if err != nil {
 			klog.Warningf("stream read error: %v", err)
 			return
@@ -123,7 +116,7 @@ func (a *AgentClient) Serve(stopCh <-chan struct{}) {
 			conn, err := net.Dial(dialReq.Protocol, dialReq.Address)
 			if err != nil {
 				resp.GetDialResponse().Error = err.Error()
-				if err := a.stream.Send(resp); err != nil {
+				if err := a.stream.RetrySend(resp); err != nil {
 					klog.Warningf("stream send error: %v", err)
 				}
 				continue
@@ -147,7 +140,7 @@ func (a *AgentClient) Serve(stopCh <-chan struct{}) {
 						resp.GetCloseResponse().Error = err.Error()
 					}
 
-					if err := a.stream.Send(resp); err != nil {
+					if err := a.stream.RetrySend(resp); err != nil {
 						klog.Warningf("close response send error: %v", err)
 					}
 
@@ -157,7 +150,7 @@ func (a *AgentClient) Serve(stopCh <-chan struct{}) {
 			}
 
 			resp.GetDialResponse().ConnectID = connID
-			if err := a.stream.Send(resp); err != nil {
+			if err := a.stream.RetrySend(resp); err != nil {
 				klog.Warningf("stream send error: %v", err)
 				continue
 			}
@@ -230,7 +223,7 @@ func (a *AgentClient) remoteToProxy(conn net.Conn, connID int64) {
 				Data:      buf[:n],
 				ConnectID: connID,
 			}}
-			if err := a.stream.Send(resp); err != nil {
+			if err := a.stream.RetrySend(resp); err != nil {
 				klog.Warningf("stream send error: %v", err)
 			}
 		}

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