Merge pull request #3 from donbing007/feature/fix-websocket-healthcheck

Feature/fix websocket healthcheck

Author: vinkdong

Gopkg.lock

@@ -6,8 +6,6 @@ import (
 	"github.com/choerodon/choerodon-cluster-agent/manager"
 	"github.com/choerodon/choerodon-cluster-agent/pkg/cluster"
 	"github.com/choerodon/choerodon-cluster-agent/pkg/cluster/kubernetes"
-	k8sclient "k8s.io/client-go/kubernetes"
-	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"github.com/choerodon/choerodon-cluster-agent/pkg/git"
 	"github.com/choerodon/choerodon-cluster-agent/pkg/helm"
 	"github.com/choerodon/choerodon-cluster-agent/pkg/kube"
@@ -17,6 +15,8 @@ import (
 	"github.com/choerodon/choerodon-cluster-agent/ws"
 	"github.com/golang/glog"
 	"github.com/spf13/pflag"
+	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	k8sclient "k8s.io/client-go/kubernetes"
 	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
 	"net/http"
 	"os"
@@ -36,12 +36,19 @@ const (
 )
 
 type AgentOptions struct {
-	Listen       string
-	UpstreamURL  string
-	Token        string
+	Listen               string
+	UpstreamURL          string
+	Token                string
+	ReadLimit            int64
+	ConnectionTimeout    time.Duration
+	WriteTimeout         time.Duration
+	HealthCheckDuration  time.Duration
+	HealthCheckTimeout   time.Duration
+	HealthCheckTryNumber int32
+
 	PrintVersion bool
 	// kubernetes controller
-	PlatformCode string
+	PlatformCode                  string
 	ConcurrentEndpointSyncs       int32
 	ConcurrentServiceSyncs        int32
 	ConcurrentRSSyncs             int32
@@ -127,8 +134,15 @@ func Run(o *AgentOptions, f cmdutil.Factory) {
 
 	checkKube(kubeClient.GetKubeClient())
 
+	appClient, err := ws.NewClient(ws.Token(o.Token), o.UpstreamURL, chans, &ws.Conf{
+		ReadLimit:            o.ReadLimit,
+		ConnectionTimeout:    o.ConnectionTimeout,
+		WriteTimeout:         o.WriteTimeout,
+		HealthCheckTryNumber: o.HealthCheckTryNumber,
+		HealthCheckTimeout:   o.HealthCheckTimeout,
+		HealthCheckDuration:  o.HealthCheckDuration,
+	})
 
-	appClient, err := ws.NewClient(ws.Token(o.Token), o.UpstreamURL, chans)
 	if err != nil {
 		errChan <- err
 		return
@@ -137,14 +151,14 @@ func Run(o *AgentOptions, f cmdutil.Factory) {
 
 	//gitRemote := git.Remote{URL: o.gitURL}
 	gitConfig := git.Config{
-		Branch:    o.gitBranch,
-		Path:      o.gitPath,
-		UserName:  o.gitUser,
-		GitUrl:    o.gitURL,
-		UserEmail: o.gitEmail,
-		SyncTag:   o.gitSyncTag,
-		DevOpsTag: o.gitDevOpsSyncTag,
-		NotesRef:  o.gitNotesRef,
+		Branch:          o.gitBranch,
+		Path:            o.gitPath,
+		UserName:        o.gitUser,
+		GitUrl:          o.gitURL,
+		UserEmail:       o.gitEmail,
+		SyncTag:         o.gitSyncTag,
+		DevOpsTag:       o.gitDevOpsSyncTag,
+		NotesRef:        o.gitNotesRef,
 		GitPollInterval: o.gitPollInterval,
 	}
 	//gitRepo := git.NewRepo(gitRemote, git.PollInterval(o.gitPollInterval))
@@ -178,7 +192,6 @@ func Run(o *AgentOptions, f cmdutil.Factory) {
 	//	k8sManifests = &kubernetes.Manifests{Namespace: o.Namespace}
 	//}
 
-
 	namespaces := manager.NewNamespaces()
 
 	ctx := controller.CreateControllerContext(
@@ -204,11 +217,11 @@ func Run(o *AgentOptions, f cmdutil.Factory) {
 			glog.Fatal(err)
 		}
 		glog.Infof("kubectl %s", kubectl)
-		cfg,_ := f.ClientConfig()
+		cfg, _ := f.ClientConfig()
 		kubectlApplier := kubernetes.NewKubectl(kubectl, cfg)
 		kubectlApplier.ApplySingleObj("kube-system", model.CRD_YAML)
 
-		k8s = kubernetes.NewCluster( kubeClient.GetKubeClient(), kubeClient.GetC7NClient(), kubectlApplier)
+		k8s = kubernetes.NewCluster(kubeClient.GetKubeClient(), kubeClient.GetC7NClient(), kubectlApplier)
 		k8sManifests = &kubernetes.Manifests{}
 	}
 	workerManager := worker.NewWorkerManager(
@@ -233,15 +246,12 @@ func Run(o *AgentOptions, f cmdutil.Factory) {
 	go workerManager.Start()
 	shutdownWg.Add(1)
 
-
-
 	go func() {
 		errChan <- http.ListenAndServe(o.Listen, nil)
 	}()
 
 }
 
-
 func (o *AgentOptions) BindFlags(fs *pflag.FlagSet) {
 	fs.BoolVar(&o.PrintVersion, "version", false, "print the version number")
 	fs.StringVar(&o.Listen, "listen", o.Listen, "address:port to listen on")
@@ -250,6 +260,12 @@ func (o *AgentOptions) BindFlags(fs *pflag.FlagSet) {
 	fs.StringVar(&o.UpstreamURL, "connect", "", "Connect to an upstream service")
 	fs.StringVar(&o.Token, "token", "", "Authentication token for upstream service")
 	fs.Int32Var(&o.ClusterId, "clusterId", 0, "the env cluster id in devops")
+	fs.Int64Var(&o.ReadLimit, "readLimit", 524288, "message read size")
+	fs.DurationVar(&o.ConnectionTimeout, "connectionTimeout", 10*time.Second, "connection server timeout, default 10 second.")
+	fs.DurationVar(&o.WriteTimeout, "writeTimeout", 3*time.Second, "write message timeout, default 3 second.")
+	fs.DurationVar(&o.HealthCheckDuration, "healthCheckDuration", 3*time.Second, "health check duration.")
+	fs.DurationVar(&o.HealthCheckTimeout, "healthCheckTimeout", 10*time.Second, "health check wait pong duration.")
+	fs.Int32Var(&o.HealthCheckTryNumber, "healthCheckTryNumber", int32(3), "health check max try number.")
 
 	// kubernetes controller
 	fs.StringVar(&o.PlatformCode, "choerodon-id", "", "choerodon platform id label")
@@ -279,13 +295,12 @@ func (o *AgentOptions) BindFlags(fs *pflag.FlagSet) {
 	fs.StringVar(&o.kubernetesKubectl, "kubernetes-kubectl", "", "Optional, explicit path to kubectl tool")
 }
 
-func checkKube(client *k8sclient.Clientset)  {
+func checkKube(client *k8sclient.Clientset) {
 	glog.Infof("check k8s role binding...")
 	_, err := client.CoreV1().Pods("").List(meta_v1.ListOptions{})
 	if err != nil {
 		glog.Errorf("check role binding failed %v", err)
-		os.Exit(0	)
+		os.Exit(0)
 	}
 	glog.Infof(" k8s role binding succeed.")
 }
-

agent/agent.go

@@ -24,6 +24,7 @@ const (
 	initialBackOff = 1 * time.Second
 	maxBackOff     = 60 * time.Second
 )
+
 var connectFlag = false
 
 type WebSocketClient interface {
@@ -43,12 +44,23 @@ type appClient struct {
 	backgroundWait sync.WaitGroup
 	pipeConns      map[string]*websocket.Conn
 	respQueue      []*model.Packet
+	conf           *Conf
+	healthCheck    *HealthCheck
+}
+
+type Conf struct {
+	ReadLimit            int64
+	ConnectionTimeout    time.Duration
+	WriteTimeout         time.Duration
+	HealthCheckDuration  time.Duration
+	HealthCheckTimeout   time.Duration
+	HealthCheckTryNumber int32
 }
 
 func NewClient(
 	t Token,
 	endpoint string,
-	crChannel *manager.CRChan) (WebSocketClient, error) {
+	crChannel *manager.CRChan, connConf *Conf) (WebSocketClient, error) {
 	if endpoint == "" {
 		return nil, fmt.Errorf("no upstream URL given")
 	}
@@ -61,13 +73,14 @@ func NewClient(
 	httpClient := cleanhttp.DefaultClient()
 
 	c := &appClient{
-		url:        endpointURL,
-		token:      t,
+		url:       endpointURL,
+		token:     t,
 		crChannel: crChannel,
-		quit:       make(chan struct{}),
-		client:     httpClient,
-		pipeConns:  make(map[string]*websocket.Conn),
-		respQueue:  make([]*model.Packet, 0, 100),
+		quit:      make(chan struct{}),
+		client:    httpClient,
+		pipeConns: make(map[string]*websocket.Conn),
+		respQueue: make([]*model.Packet, 0, 100),
+		conf:      connConf,
 	}
 
 	return c, nil
@@ -102,12 +115,18 @@ func (c *appClient) Loop(stop <-chan struct{}, done *sync.WaitGroup) {
 func (c *appClient) connect() error {
 	glog.V(1).Info("Start connect to DevOps service")
 	var err error
-	c.conn, err = dial(c.url.String(), c.token)
+	c.conn, err = dial(c.url.String(), c.token, c.conf.ConnectionTimeout)
 	if err != nil {
 		return err
 	}
 	glog.V(1).Info("Connect to DevOps service success")
 
+	// 配置连接.
+	c.configure()
+
+	// 开始健康检查
+	c.startHealthCheck()
+
 	// 建立连接，同步资源对象
 	if connectFlag {
 		c.crChannel.CommandChan <- newReConnectCommand()
@@ -124,7 +143,6 @@ func (c *appClient) connect() error {
 	go func() {
 		defer close(done)
 
-		c.conn.SetPingHandler(nil)
 		for {
 			var command model.Packet
 			err := c.conn.ReadJSON(&command)
@@ -134,14 +152,17 @@ func (c *appClient) connect() error {
 				}
 				break
 			}
+
+			c.healthCheck.OnRecvice(nil)
+
 			glog.V(1).Info("receive command: ", command)
 			c.crChannel.CommandChan <- &command
 		}
 	}()
 
 	end := 0
 	for ; end < len(c.respQueue); end++ {
-		c.conn.SetWriteDeadline(time.Now().Add(WriteWait))
+		c.conn.SetWriteDeadline(time.Now().Add(c.conf.WriteTimeout))
 		resp := c.respQueue[end]
 		if err := c.sendResponse(resp); err != nil {
 			c.respQueue = c.respQueue[end:]
@@ -292,7 +313,7 @@ func (c *appClient) pipeConnection(id string, pipe common.Pipe) (bool, error) {
 	}
 	newURLStr := fmt.Sprintf("%s.%s:%s", newURL.String(), pipe.PipeType(), id)
 	headers := http.Header{}
-	conn, resp, err := dialWS(newURLStr, headers)
+	conn, resp, err := dialWS(newURLStr, headers, c.conf.ConnectionTimeout)
 	if resp != nil && resp.StatusCode == http.StatusNotFound {
 		glog.V(2).Info("response with not found")
 		pipe.Close()
@@ -317,16 +338,55 @@ func (c *appClient) pipeConnection(id string, pipe common.Pipe) (bool, error) {
 	return true, nil
 }
 
-func newReConnectCommand()  *model.Packet{
+func (c *appClient) configure() {
+	c.conn.SetReadLimit(c.conf.ReadLimit)
+
+	c.conn.SetPingHandler(func(message string) error {
+		return c.healthCheck.OnPing([]byte(message))
+	})
+
+	c.conn.SetPongHandler(func(message string) error {
+		return c.healthCheck.OnPone([]byte(message))
+	})
+
+	c.conn.SetCloseHandler(func(code int, text string) error {
+		c.healthCheck.Stop()
+
+		return nil
+	})
+}
+
+func (c *appClient) startHealthCheck() error {
+
+	if c.healthCheck != nil {
+		c.healthCheck = nil
+	}
+
+	h, err := NewHealthCheck(c.conf, c.conn, func() {
+		e := c.conn.Close()
+		if e != nil {
+			glog.V(1).Infof("An error occurred while closing target %s because of %v.", c.conn.RemoteAddr().String(), e)
+		}
+	})
+
+	if err != nil {
+		return err
+	} else {
+		c.healthCheck = h
+		return nil
+	}
+}
+
+func newReConnectCommand() *model.Packet {
 	return &model.Packet{
-		Key: "inter:inter",
+		Key:  "inter:inter",
 		Type: model.ReSyncAgent,
 	}
 }
-func newUpgradeInfoCommand(connectUrl string)  *model.Packet{
+func newUpgradeInfoCommand(connectUrl string) *model.Packet {
 	return &model.Packet{
-		Key: "inter:inter",
-		Type: model.UpgradeCluster,
+		Key:     "inter:inter",
+		Type:    model.UpgradeCluster,
 		Payload: connectUrl,
 	}
 }

ws/client.go

@@ -1,5 +1,6 @@
 package ws
 
+
 import (
 	"encoding/json"
 	"fmt"