Merge pull request fatedier#998 from fatedier/health

frpc: support health check

Makefile

@@ -34,7 +34,7 @@ gotest:
 	go test -v --cover ./utils/...
 
 ci:
-	go test -count=1 -v ./tests/...
+	go test -count=1 -p=1 -v ./tests/...
 
 alltest: gotest ci
 

client/admin_api.go

@@ -24,6 +24,7 @@ import (
 
 	ini "github.com/vaughan0/go-ini"
 
+	"github.com/fatedier/frp/client/proxy"
 	"github.com/fatedier/frp/g"
 	"github.com/fatedier/frp/models/config"
 	"github.com/fatedier/frp/utils/log"
@@ -121,7 +122,7 @@ func (a ByProxyStatusResp) Len() int           { return len(a) }
 func (a ByProxyStatusResp) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
 func (a ByProxyStatusResp) Less(i, j int) bool { return strings.Compare(a[i].Name, a[j].Name) < 0 }
 
-func NewProxyStatusResp(status *ProxyStatus) ProxyStatusResp {
+func NewProxyStatusResp(status *proxy.ProxyStatus) ProxyStatusResp {
 	psr := ProxyStatusResp{
 		Name:   status.Name,
 		Type:   status.Type,

client/control.go

@@ -21,6 +21,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/fatedier/frp/client/proxy"
 	"github.com/fatedier/frp/g"
 	"github.com/fatedier/frp/models/config"
 	"github.com/fatedier/frp/models/msg"
@@ -37,7 +38,8 @@ type Control struct {
 	runId string
 
 	// manage all proxies
-	pm *ProxyManager
+	pxyCfgs map[string]config.ProxyConf
+	pm      *proxy.ProxyManager
 
 	// manage all visitors
 	vm *VisitorManager
@@ -76,6 +78,7 @@ func NewControl(runId string, conn frpNet.Conn, session *fmux.Session, pxyCfgs m
 		runId:              runId,
 		conn:               conn,
 		session:            session,
+		pxyCfgs:            pxyCfgs,
 		sendCh:             make(chan msg.Message, 100),
 		readCh:             make(chan msg.Message, 100),
 		closedCh:           make(chan struct{}),
@@ -85,8 +88,8 @@ func NewControl(runId string, conn frpNet.Conn, session *fmux.Session, pxyCfgs m
 		msgHandlerShutdown: shutdown.New(),
 		Logger:             log.NewPrefixLogger(""),
 	}
-	ctl.pm = NewProxyManager(ctl.sendCh, "")
-	ctl.pm.Reload(pxyCfgs, false)
+	ctl.pm = proxy.NewProxyManager(ctl.sendCh, runId)
+
 	ctl.vm = NewVisitorManager(ctl)
 	ctl.vm.Reload(visitorCfgs)
 	return ctl
@@ -95,10 +98,10 @@ func NewControl(runId string, conn frpNet.Conn, session *fmux.Session, pxyCfgs m
 func (ctl *Control) Run() {
 	go ctl.worker()
 
-	// start all local visitors and send NewProxy message for all configured proxies
-	ctl.pm.Reset(ctl.sendCh, ctl.runId)
-	ctl.pm.CheckAndStartProxy([]string{ProxyStatusNew})
+	// start all proxies
+	ctl.pm.Reload(ctl.pxyCfgs)
 
+	// start all visitors
 	go ctl.vm.Run()
 	return
 }
@@ -142,7 +145,7 @@ func (ctl *Control) HandleNewProxyResp(inMsg *msg.NewProxyResp) {
 }
 
 func (ctl *Control) Close() error {
-	ctl.pm.CloseProxies()
+	ctl.conn.Close()
 	return nil
 }
 
@@ -275,33 +278,26 @@ func (ctl *Control) worker() {
 	go ctl.reader()
 	go ctl.writer()
 
-	checkInterval := 60 * time.Second
-	checkProxyTicker := time.NewTicker(checkInterval)
-
-	for {
-		select {
-		case <-checkProxyTicker.C:
-			// check which proxy registered failed and reregister it to server
-			ctl.pm.CheckAndStartProxy([]string{ProxyStatusStartErr, ProxyStatusClosed})
-		case <-ctl.closedCh:
-			// close related channels and wait until other goroutines done
-			close(ctl.readCh)
-			ctl.readerShutdown.WaitDone()
-			ctl.msgHandlerShutdown.WaitDone()
+	select {
+	case <-ctl.closedCh:
+		// close related channels and wait until other goroutines done
+		close(ctl.readCh)
+		ctl.readerShutdown.WaitDone()
+		ctl.msgHandlerShutdown.WaitDone()
 
-			close(ctl.sendCh)
-			ctl.writerShutdown.WaitDone()
+		close(ctl.sendCh)
+		ctl.writerShutdown.WaitDone()
 
-			ctl.pm.CloseProxies()
+		ctl.pm.Close()
+		ctl.vm.Close()
 
-			close(ctl.closedDoneCh)
-			return
-		}
+		close(ctl.closedDoneCh)
+		return
 	}
 }
 
 func (ctl *Control) ReloadConf(pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.VisitorConf) error {
 	ctl.vm.Reload(visitorCfgs)
-	err := ctl.pm.Reload(pxyCfgs, true)
-	return err
+	ctl.pm.Reload(pxyCfgs)
+	return nil
 }

client/event/event.go

@@ -0,0 +1,28 @@
+package event
+
+import (
+	"errors"
+
+	"github.com/fatedier/frp/models/msg"
+)
+
+type EventType int
+
+const (
+	EvStartProxy EventType = iota
+	EvCloseProxy
+)
+
+var (
+	ErrPayloadType = errors.New("error payload type")
+)
+
+type EventHandler func(evType EventType, payload interface{}) error
+
+type StartProxyPayload struct {
+	NewProxyMsg *msg.NewProxy
+}
+
+type CloseProxyPayload struct {
+	CloseProxyMsg *msg.CloseProxy
+}

client/health.go → client/health/health.go

@@ -12,13 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package client
+package health
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"net"
 	"net/http"
 	"time"
+
+	"github.com/fatedier/frp/utils/log"
+)
+
+var (
+	ErrHealthCheckType = errors.New("error health check type")
 )
 
 type HealthCheckMonitor struct {
@@ -40,6 +48,8 @@ type HealthCheckMonitor struct {
 
 	ctx    context.Context
 	cancel context.CancelFunc
+
+	l log.Logger
 }
 
 func NewHealthCheckMonitor(checkType string, intervalS int, timeoutS int, maxFailedTimes int, addr string, url string,
@@ -70,6 +80,10 @@ func NewHealthCheckMonitor(checkType string, intervalS int, timeoutS int, maxFai
 	}
 }
 
+func (monitor *HealthCheckMonitor) SetLogger(l log.Logger) {
+	monitor.l = l
+}
+
 func (monitor *HealthCheckMonitor) Start() {
 	go monitor.checkWorker()
 }
@@ -81,7 +95,7 @@ func (monitor *HealthCheckMonitor) Stop() {
 func (monitor *HealthCheckMonitor) checkWorker() {
 	for {
 		ctx, cancel := context.WithDeadline(monitor.ctx, time.Now().Add(monitor.timeout))
-		ok := monitor.doCheck(ctx)
+		err := monitor.doCheck(ctx)
 
 		// check if this monitor has been closed
 		select {
@@ -92,14 +106,26 @@ func (monitor *HealthCheckMonitor) checkWorker() {
 			cancel()
 		}
 
-		if ok {
+		if err == nil {
+			if monitor.l != nil {
+				monitor.l.Trace("do one health check success")
+			}
 			if !monitor.statusOK && monitor.statusNormalFn != nil {
+				if monitor.l != nil {
+					monitor.l.Info("health check status change to success")
+				}
 				monitor.statusOK = true
 				monitor.statusNormalFn()
 			}
 		} else {
+			if monitor.l != nil {
+				monitor.l.Warn("do one health check failed: %v", err)
+			}
 			monitor.failedTimes++
 			if monitor.statusOK && int(monitor.failedTimes) >= monitor.maxFailedTimes && monitor.statusFailedFn != nil {
+				if monitor.l != nil {
+					monitor.l.Warn("health check status change to failed")
+				}
 				monitor.statusOK = false
 				monitor.statusFailedFn()
 			}
@@ -109,39 +135,44 @@ func (monitor *HealthCheckMonitor) checkWorker() {
 	}
 }
 
-func (monitor *HealthCheckMonitor) doCheck(ctx context.Context) bool {
+func (monitor *HealthCheckMonitor) doCheck(ctx context.Context) error {
 	switch monitor.checkType {
 	case "tcp":
 		return monitor.doTcpCheck(ctx)
 	case "http":
 		return monitor.doHttpCheck(ctx)
 	default:
-		return false
+		return ErrHealthCheckType
 	}
 }
 
-func (monitor *HealthCheckMonitor) doTcpCheck(ctx context.Context) bool {
+func (monitor *HealthCheckMonitor) doTcpCheck(ctx context.Context) error {
+	// if tcp address is not specified, always return nil
+	if monitor.addr == "" {
+		return nil
+	}
+
 	var d net.Dialer
 	conn, err := d.DialContext(ctx, "tcp", monitor.addr)
 	if err != nil {
-		return false
+		return err
 	}
 	conn.Close()
-	return true
+	return nil
 }
 
-func (monitor *HealthCheckMonitor) doHttpCheck(ctx context.Context) bool {
+func (monitor *HealthCheckMonitor) doHttpCheck(ctx context.Context) error {
 	req, err := http.NewRequest("GET", monitor.url, nil)
 	if err != nil {
-		return false
+		return err
 	}
 	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
-		return false
+		return err
 	}
 
 	if resp.StatusCode/100 != 2 {
-		return false
+		return fmt.Errorf("do http health check, StatusCode is [%d] not 2xx", resp.StatusCode)
 	}
-	return true
+	return nil
 }

client/proxy.go → client/proxy/proxy.go

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package client
+package proxy
 
 import (
 	"bytes"