健康检查

client/client.go

@@ -2,6 +2,7 @@ package client
 
 import (
 	"github.com/cnlh/nps/lib/common"
+	"github.com/cnlh/nps/lib/config"
 	"github.com/cnlh/nps/lib/conn"
 	"github.com/cnlh/nps/lib/mux"
 	"github.com/cnlh/nps/vender/github.com/astaxie/beego/logs"
@@ -19,16 +20,18 @@ type TRPClient struct {
 	vKey           string
 	tunnel         *mux.Mux
 	signal         *conn.Conn
+	cnf            *config.Config
 }
 
 //new client
-func NewRPClient(svraddr string, vKey string, bridgeConnType string, proxyUrl string) *TRPClient {
+func NewRPClient(svraddr string, vKey string, bridgeConnType string, proxyUrl string, cnf *config.Config) *TRPClient {
 	return &TRPClient{
 		svrAddr:        svraddr,
 		vKey:           vKey,
 		bridgeConnType: bridgeConnType,
 		stop:           make(chan bool),
 		proxyUrl:       proxyUrl,
+		cnf:            cnf,
 	}
 }
 
@@ -55,6 +58,7 @@ func (s *TRPClient) Close() {
 func (s *TRPClient) processor(c *conn.Conn) {
 	s.signal = c
 	go s.dealChan()
+	go heathCheck(s.cnf, c)
 	for {
 		flags, err := c.ReadFlag()
 		if err != nil {

client/health.go

@@ -0,0 +1,104 @@
+package client
+
+import (
+	"container/heap"
+	"github.com/cnlh/nps/lib/config"
+	"github.com/cnlh/nps/lib/file"
+	"github.com/cnlh/nps/lib/sheap"
+	"net"
+	"net/http"
+	"strings"
+	"time"
+)
+
+func heathCheck(cnf *config.Config, c net.Conn) {
+	var hosts []*file.Host
+	var tunnels []*file.Tunnel
+	h := &sheap.IntHeap{}
+	for _, v := range cnf.Hosts {
+		if v.HealthMaxFail > 0 && v.HealthCheckTimeout > 0 && v.HealthCheckInterval > 0 {
+			v.HealthNextTime = time.Now().Add(time.Duration(v.HealthCheckInterval))
+			heap.Push(h, v.HealthNextTime.Unix())
+			v.HealthMap = make(map[string]int)
+			hosts = append(hosts, v)
+		}
+	}
+	for _, v := range cnf.Tasks {
+		if v.HealthMaxFail > 0 && v.HealthCheckTimeout > 0 && v.HealthCheckInterval > 0 {
+			v.HealthNextTime = time.Now().Add(time.Duration(v.HealthCheckInterval))
+			heap.Push(h, v.HealthNextTime.Unix())
+			v.HealthMap = make(map[string]int)
+			tunnels = append(tunnels, v)
+		}
+	}
+	if len(hosts) == 0 && len(tunnels) == 0 {
+		return
+	}
+	for {
+		rs := heap.Pop(h).(int64) - time.Now().Unix()
+		if rs < 0 {
+			continue
+		}
+		timer := time.NewTicker(time.Duration(rs))
+		select {
+		case <-timer.C:
+			for _, v := range hosts {
+				if v.HealthNextTime.Before(time.Now()) {
+					v.HealthNextTime = time.Now().Add(time.Duration(v.HealthCheckInterval))
+					//check
+					go checkHttp(v, c)
+					//reset time
+					heap.Push(h, v.HealthNextTime.Unix())
+				}
+			}
+			for _, v := range tunnels {
+				if v.HealthNextTime.Before(time.Now()) {
+					v.HealthNextTime = time.Now().Add(time.Duration(v.HealthCheckInterval))
+					//check
+					go checkTcp(v, c)
+					//reset time
+					heap.Push(h, v.HealthNextTime.Unix())
+				}
+			}
+		}
+	}
+}
+
+func checkTcp(t *file.Tunnel, c net.Conn) {
+	arr := strings.Split(t.Target, "\n")
+	for _, v := range arr {
+		if _, err := net.DialTimeout("tcp", v, time.Duration(t.HealthCheckTimeout)); err != nil {
+			t.HealthMap[v] += 1
+		}
+		if t.HealthMap[v] > t.HealthMaxFail {
+			t.HealthMap[v] += 1
+			if t.HealthMap[v] == t.HealthMaxFail {
+				//send fail remove
+				ch <- file.NewHealthInfo("tcp", v, true)
+			}
+		} else if t.HealthMap[v] >= t.HealthMaxFail {
+			//send recovery add
+			ch <- file.NewHealthInfo("tcp", v, false)
+			t.HealthMap[v] = 0
+		}
+	}
+}
+
+func checkHttp(h *file.Host, ch chan *file.HealthInfo) {
+	arr := strings.Split(h.Target, "\n")
+	client := &http.Client{}
+	client.Timeout = time.Duration(h.HealthCheckTimeout) * time.Second
+	for _, v := range arr {
+		if _, err := client.Get(v + h.HttpHealthUrl); err != nil {
+			h.HealthMap[v] += 1
+			if h.HealthMap[v] == h.HealthMaxFail {
+				//send fail remove
+				ch <- file.NewHealthInfo("http", v, true)
+			}
+		} else if h.HealthMap[v] >= h.HealthMaxFail {
+			//send recovery add
+			h.HealthMap[v] = 0
+			ch <- file.NewHealthInfo("http", v, false)
+		}
+	}
+}