Merge pull request #46 from caesarxuchao/agent-connect-to-all-servers

Make proxy work even if proxy server is HA and is behind a LB

Author: k8s-ci-robot

cmd/agent/main.go

@@ -22,7 +22,9 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"time"
 
+	"github.com/google/uuid"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
@@ -61,6 +63,22 @@ type GrpcProxyAgentOptions struct {
 	// Configuration for connecting to the proxy-server
 	proxyServerHost string
 	proxyServerPort 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 {
@@ -70,6 +88,10 @@ func (o *GrpcProxyAgentOptions) Flags() *pflag.FlagSet {
 	flags.StringVar(&o.caCert, "ca-cert", o.caCert, "If non-empty the CAs we use to validate clients.")
 	flags.StringVar(&o.proxyServerHost, "proxy-server-host", o.proxyServerHost, "The hostname to use to connect to the proxy-server.")
 	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
 }
 
@@ -79,6 +101,10 @@ func (o *GrpcProxyAgentOptions) Print() {
 	klog.Warningf("CACert set to \"%s\".\n", o.caCert)
 	klog.Warningf("ProxyServerHost set to \"%s\".\n", o.proxyServerHost)
 	klog.Warningf("ProxyServerPort set to %d.\n", o.proxyServerPort)
+	klog.Warningf("AgentID set to %s.\n", o.agentID)
+	klog.Warningf("SyncInterval set to %d seconds.\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 {
@@ -111,11 +137,15 @@ func (o *GrpcProxyAgentOptions) Validate() error {
 
 func newGrpcProxyAgentOptions() *GrpcProxyAgentOptions {
 	o := GrpcProxyAgentOptions{
-		agentCert:       "",
-		agentKey:        "",
-		caCert:          "",
-		proxyServerHost: "127.0.0.1",
-		proxyServerPort: 8091,
+		agentCert:         "",
+		agentKey:          "",
+		caCert:            "",
+		proxyServerHost:   "127.0.0.1",
+		proxyServerPort:   8091,
+		agentID:           uuid.New().String(),
+		syncInterval:      5,
+		probeInterval:     5,
+		reconnectInterval: 5,
 	}
 	return &o
 }
@@ -162,14 +192,9 @@ func (a *Agent) runProxyConnection(o *GrpcProxyAgentOptions) error {
 		return err
 	}
 	dialOption := grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig))
-	client, err := agentclient.NewAgentClient(fmt.Sprintf("%s:%d", o.proxyServerHost, o.proxyServerPort), dialOption)
-	if err != nil {
-		return err
-	}
-
-	stopCh := make(chan struct{})
-
-	go client.Serve(stopCh)
+	cc := o.ClientSetConfig(dialOption)
+	cs := cc.NewAgentClientSet()
+	cs.Serve()
 
 	return nil
 }

cmd/proxy/main.go

@@ -23,13 +23,15 @@ import (
 	"flag"
 	"fmt"
 	"io/ioutil"
-	"k8s.io/klog"
 	"net"
 	"net/http"
 	"os"
 	"os/signal"
 	"syscall"
 
+	"k8s.io/klog"
+
+	"github.com/google/uuid"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
@@ -79,6 +81,11 @@ type ProxyRunOptions struct {
 	agentPort uint
 	// Port we listen for admin connections on.
 	adminPort uint
+
+	// ID of this server.
+	serverID string
+	// Number of proxy server instances, should be 1 unless it is a HA server.
+	serverCount uint
 }
 
 func (o *ProxyRunOptions) Flags() *pflag.FlagSet {
@@ -94,6 +101,8 @@ func (o *ProxyRunOptions) Flags() *pflag.FlagSet {
 	flags.UintVar(&o.serverPort, "server-port", o.serverPort, "Port we listen for server connections on. Set to 0 for UDS.")
 	flags.UintVar(&o.agentPort, "agent-port", o.agentPort, "Port we listen for agent connections on.")
 	flags.UintVar(&o.adminPort, "admin-port", o.adminPort, "Port we listen for admin connections on.")
+	flags.StringVar(&o.serverID, "server-id", o.serverID, "The unique ID of this server.")
+	flags.UintVar(&o.serverCount, "server-count", o.serverCount, "The number of proxy server instances, should be 1 unless it is an HA server.")
 	return flags
 }
 
@@ -109,6 +118,8 @@ func (o *ProxyRunOptions) Print() {
 	klog.Warningf("Server port set to %d.\n", o.serverPort)
 	klog.Warningf("Agent port set to %d.\n", o.agentPort)
 	klog.Warningf("Admin port set to %d.\n", o.adminPort)
+	klog.Warningf("ServerID set to %s.\n", o.serverID)
+	klog.Warningf("ServerCount set to %d.\n", o.serverCount)
 }
 
 func (o *ProxyRunOptions) Validate() error {
@@ -207,6 +218,8 @@ func newProxyRunOptions() *ProxyRunOptions {
 		serverPort:    8090,
 		agentPort:     8091,
 		adminPort:     8092,
+		serverID:      uuid.New().String(),
+		serverCount:   1,
 	}
 	return &o
 }
@@ -234,7 +247,7 @@ func (p *Proxy) run(o *ProxyRunOptions) error {
 		return fmt.Errorf("failed to validate server options with %v", err)
 	}
 	ctx, cancel := context.WithCancel(context.Background())
-	server := agentserver.NewProxyServer()
+	server := agentserver.NewProxyServer(o.serverID, int(o.serverCount))
 
 	klog.Info("Starting master server for client connections.")
 	masterStop, err := p.runMasterServer(ctx, o, server)

go.mod

@@ -5,13 +5,17 @@ go 1.12
 require (
 	github.com/beorn7/perks v1.0.0 // indirect
 	github.com/golang/protobuf v1.3.2
+	github.com/google/uuid v1.1.1
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect
 	github.com/prometheus/client_golang v0.9.2
 	github.com/prometheus/common v0.4.0 // indirect
 	github.com/prometheus/procfs v0.0.0-20190507164030-5867b95ac084 // indirect
 	github.com/spf13/cobra v0.0.3
 	github.com/spf13/pflag v1.0.3
+	golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3 // indirect
 	golang.org/x/sys v0.0.0-20190225065934-cc5685c2db12 // indirect
+	golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135 // indirect
 	google.golang.org/grpc v1.26.0
+	honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc // indirect
 	k8s.io/klog v0.3.0
 )

go.sum

@@ -21,9 +21,9 @@ github.com/golang/protobuf v1.2.0 h1:P3YflyNX/ehuJFLhxviNdFxQPkGK5cDcApsge1SqnvM
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.2 h1:6nsPYzhq5kReh6QImI3k5qWzO4PEbvbIW2cwSfR/6xs=
 github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
-github.com/google/go-cmp v0.2.0 h1:+dTQ8DZQJz0Mb/HjFlkptS1FeQ4cWSnN941F8aEG4SQ=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
-github.com/inconshreveable/mousetrap v1.0.0 h1:Z8tu5sraLXCXIcARxBp/8cbvlwVa7Z1NHg9XEKhtSvM=
+github.com/google/uuid v1.1.1 h1:Gkbcsh/GbpXz7lPftLA3P6TYMwjCLYm83jiFQZF/3gY=
+github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
 github.com/julienschmidt/httprouter v1.2.0/go.mod h1:SYymIcj16QtmaHHD7aYtjjsJG7VTCxuUUipMqKk8s4w=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=

pkg/agent/agentclient/client.go

@@ -34,20 +34,32 @@ type AgentClient struct {
 	connContext map[int64]*connContext
 
 	stream *RedialableAgentClient
+	stopCh <-chan struct{}
 }
 
-// NewAgentClient creates an AgentClient
-func NewAgentClient(address string, opts ...grpc.DialOption) (*AgentClient, error) {
-	stream, err := NewRedialableAgentClient(address, opts...)
+func newAgentClient(address, agentID string, cs *ClientSet, opts ...grpc.DialOption) (*AgentClient, error) {
+	stream, err := NewRedialableAgentClient(address, agentID, cs, opts...)
 	if err != nil {
 		return nil, err
 	}
+	return newAgentClientWithRedialableAgentClient(stream), nil
+}
 
-	a := &AgentClient{
+func newAgentClientWithRedialableAgentClient(rac *RedialableAgentClient) *AgentClient {
+	return &AgentClient{
 		connContext: make(map[int64]*connContext),
-		stream:      stream,
+		stream:      rac,
+		stopCh:      rac.stopCh,
+	}
+}
+
+// Close closes the underlying stream.
+func (c *AgentClient) Close() {
+	if c.stream == nil {
+		klog.Warning("Unexpected empty AgentClient.stream")
+		return
 	}
-	return a, nil
+	c.stream.Close()
 }
 
 // connContext tracks a connection from agent to node network.
@@ -75,10 +87,12 @@ func (c *connContext) cleanup() {
 // gRPC stream. Successful Connect is required before Serve. The
 // The requests include things like opening a connection to a server,
 // streaming data and close the connection.
-func (a *AgentClient) Serve(stopCh <-chan struct{}) {
+func (a *AgentClient) Serve() {
+	klog.Infof("Start serving for serverID %s", a.stream.serverID)
+	go a.stream.probe()
 	for {
 		select {
-		case <-stopCh:
+		case <-a.stopCh:
 			klog.Info("stop agent client.")
 			return
 		default:

pkg/agent/agentclient/client_test.go

@@ -17,14 +17,15 @@ import (
 func TestServeData_HTTP(t *testing.T) {
 	var err error
 	var stream agent.AgentService_ConnectClient
+	stopCh := make(chan struct{})
 	client := &AgentClient{
 		connContext: make(map[int64]*connContext),
+		stopCh:      stopCh,
 	}
 	client.stream, stream = pipe2()
-	stopCh := make(chan struct{})
 
 	// Start agent
-	go client.Serve(stopCh)
+	go client.Serve()
 	defer close(stopCh)
 
 	// Start test http server as remote service
@@ -112,14 +113,15 @@ func TestServeData_HTTP(t *testing.T) {
 
 func TestClose_Client(t *testing.T) {
 	var stream agent.AgentService_ConnectClient
+	stopCh := make(chan struct{})
 	client := &AgentClient{
 		connContext: make(map[int64]*connContext),
+		stopCh:      stopCh,
 	}
 	client.stream, stream = pipe2()
-	stopCh := make(chan struct{})
 
 	// Start agent
-	go client.Serve(stopCh)
+	go client.Serve()
 	defer close(stopCh)
 
 	// Start test http server as remote service