frp/client/health/health.go

178 lines
4.0 KiB
Go
Raw Normal View History

// Copyright 2018 fatedier, fatedier@gmail.com
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
2018-12-09 22:06:22 +08:00
package health
import (
2018-07-16 01:21:29 +08:00
"context"
2018-12-07 17:05:36 +08:00
"errors"
"fmt"
"io"
2018-07-16 01:21:29 +08:00
"net"
"net/http"
"strings"
2018-07-16 01:21:29 +08:00
"time"
2018-12-07 17:05:36 +08:00
v1 "github.com/fatedier/frp/pkg/config/v1"
2020-09-23 13:49:14 +08:00
"github.com/fatedier/frp/pkg/util/xlog"
2018-12-07 17:05:36 +08:00
)
2022-08-29 01:02:53 +08:00
var ErrHealthCheckType = errors.New("error health check type")
2020-05-24 17:48:37 +08:00
type Monitor struct {
2018-07-16 01:21:29 +08:00
checkType string
interval time.Duration
timeout time.Duration
maxFailedTimes int
// For tcp
addr string
// For http
url string
failedTimes uint64
statusOK bool
statusNormalFn func()
statusFailedFn func()
ctx context.Context
cancel context.CancelFunc
}
func NewMonitor(ctx context.Context, cfg v1.HealthCheckConfig, addr string,
2022-08-29 01:02:53 +08:00
statusNormalFn func(), statusFailedFn func(),
) *Monitor {
if cfg.IntervalSeconds <= 0 {
cfg.IntervalSeconds = 10
2018-07-16 01:21:29 +08:00
}
if cfg.TimeoutSeconds <= 0 {
cfg.TimeoutSeconds = 3
2018-07-16 01:21:29 +08:00
}
if cfg.MaxFailed <= 0 {
cfg.MaxFailed = 1
2018-07-16 01:21:29 +08:00
}
2019-10-12 20:13:12 +08:00
newctx, cancel := context.WithCancel(ctx)
var url string
if cfg.Type == "http" && cfg.Path != "" {
s := "http://" + addr
if !strings.HasPrefix(cfg.Path, "/") {
s += "/"
}
url = s + cfg.Path
}
2020-05-24 17:48:37 +08:00
return &Monitor{
checkType: cfg.Type,
interval: time.Duration(cfg.IntervalSeconds) * time.Second,
timeout: time.Duration(cfg.TimeoutSeconds) * time.Second,
maxFailedTimes: cfg.MaxFailed,
2018-07-16 01:21:29 +08:00
addr: addr,
url: url,
statusOK: false,
statusNormalFn: statusNormalFn,
statusFailedFn: statusFailedFn,
2019-10-12 20:13:12 +08:00
ctx: newctx,
2018-07-16 01:21:29 +08:00
cancel: cancel,
}
}
2020-05-24 17:48:37 +08:00
func (monitor *Monitor) Start() {
2018-07-16 01:21:29 +08:00
go monitor.checkWorker()
}
2020-05-24 17:48:37 +08:00
func (monitor *Monitor) Stop() {
2018-07-16 01:21:29 +08:00
monitor.cancel()
}
2020-05-24 17:48:37 +08:00
func (monitor *Monitor) checkWorker() {
2019-10-12 20:13:12 +08:00
xl := xlog.FromContextSafe(monitor.ctx)
2018-07-16 01:21:29 +08:00
for {
2019-08-09 12:47:27 +08:00
doCtx, cancel := context.WithDeadline(monitor.ctx, time.Now().Add(monitor.timeout))
err := monitor.doCheck(doCtx)
2018-07-16 01:21:29 +08:00
// check if this monitor has been closed
select {
2019-08-09 12:47:27 +08:00
case <-monitor.ctx.Done():
2018-07-16 01:21:29 +08:00
cancel()
return
default:
cancel()
}
2018-12-07 17:05:36 +08:00
if err == nil {
2024-03-12 13:58:53 +08:00
xl.Tracef("do one health check success")
2018-07-16 01:21:29 +08:00
if !monitor.statusOK && monitor.statusNormalFn != nil {
2024-03-12 13:58:53 +08:00
xl.Infof("health check status change to success")
2018-07-16 01:21:29 +08:00
monitor.statusOK = true
monitor.statusNormalFn()
}
} else {
2024-03-12 13:58:53 +08:00
xl.Warnf("do one health check failed: %v", err)
2018-07-16 01:21:29 +08:00
monitor.failedTimes++
if monitor.statusOK && int(monitor.failedTimes) >= monitor.maxFailedTimes && monitor.statusFailedFn != nil {
2024-03-12 13:58:53 +08:00
xl.Warnf("health check status change to failed")
2018-07-16 01:21:29 +08:00
monitor.statusOK = false
monitor.statusFailedFn()
}
}
time.Sleep(monitor.interval)
}
}
2020-05-24 17:48:37 +08:00
func (monitor *Monitor) doCheck(ctx context.Context) error {
2018-07-16 01:21:29 +08:00
switch monitor.checkType {
case "tcp":
2020-05-24 17:48:37 +08:00
return monitor.doTCPCheck(ctx)
2018-07-16 01:21:29 +08:00
case "http":
2020-05-24 17:48:37 +08:00
return monitor.doHTTPCheck(ctx)
2018-07-16 01:21:29 +08:00
default:
2018-12-07 17:05:36 +08:00
return ErrHealthCheckType
2018-07-16 01:21:29 +08:00
}
}
2020-05-24 17:48:37 +08:00
func (monitor *Monitor) doTCPCheck(ctx context.Context) error {
2018-12-07 17:05:36 +08:00
// if tcp address is not specified, always return nil
if monitor.addr == "" {
return nil
}
2018-07-16 01:21:29 +08:00
var d net.Dialer
conn, err := d.DialContext(ctx, "tcp", monitor.addr)
if err != nil {
2018-12-07 17:05:36 +08:00
return err
2018-07-16 01:21:29 +08:00
}
conn.Close()
2018-12-07 17:05:36 +08:00
return nil
2018-07-16 01:21:29 +08:00
}
2020-05-24 17:48:37 +08:00
func (monitor *Monitor) doHTTPCheck(ctx context.Context) error {
2022-08-29 01:02:53 +08:00
req, err := http.NewRequestWithContext(ctx, "GET", monitor.url, nil)
2018-07-16 01:21:29 +08:00
if err != nil {
2018-12-07 17:05:36 +08:00
return err
2018-07-16 01:21:29 +08:00
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
2018-12-07 17:05:36 +08:00
return err
2018-07-16 01:21:29 +08:00
}
defer resp.Body.Close()
2022-08-29 01:02:53 +08:00
_, _ = io.Copy(io.Discard, resp.Body)
2018-07-16 01:21:29 +08:00
if resp.StatusCode/100 != 2 {
2018-12-07 17:05:36 +08:00
return fmt.Errorf("do http health check, StatusCode is [%d] not 2xx", resp.StatusCode)
2018-07-16 01:21:29 +08:00
}
2018-12-07 17:05:36 +08:00
return nil
}