goproxy/internal/healthcheck/healthchecker.go

package healthcheck

import (
	"context"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"
)

// HealthChecker 健康检查器
type HealthChecker struct {
	// 配置
	config *Config
	// 健康检查状态
	statusMap sync.Map
	// 状态变更回调
	statusChangeCallback func(string, bool)
	// 是否运行中
	running bool
	// 上下文
	ctx context.Context
	// 取消函数
	cancel context.CancelFunc
	// 互斥锁
	mu sync.Mutex
}

// Config 健康检查配置
type Config struct {
	// 检查间隔
	Interval time.Duration
	// 检查超时
	Timeout time.Duration
	// 检查路径
	Path string
	// 检查方法
	Method string
	// 检查状态码
	SuccessStatus int
	// 最大失败次数
	MaxFails int
	// 最小成功次数
	MinSuccess int
}

// NewHealthChecker 创建健康检查器
func NewHealthChecker(config *Config) *HealthChecker {
	if config.Path == "" {
		config.Path = "/"
	}
	if config.Method == "" {
		config.Method = http.MethodGet
	}
	if config.SuccessStatus == 0 {
		config.SuccessStatus = http.StatusOK
	}
	if config.MaxFails == 0 {
		config.MaxFails = 3
	}
	if config.MinSuccess == 0 {
		config.MinSuccess = 2
	}

	ctx, cancel := context.WithCancel(context.Background())
	return &HealthChecker{
		config:  config,
		ctx:     ctx,
		cancel:  cancel,
		running: false,
	}
}

// Start 启动健康检查
func (hc *HealthChecker) Start() {
	hc.mu.Lock()
	defer hc.mu.Unlock()

	if hc.running {
		return
	}

	hc.running = true
	go hc.run()
}

// Stop 停止健康检查
func (hc *HealthChecker) Stop() {
	hc.mu.Lock()
	defer hc.mu.Unlock()

	if !hc.running {
		return
	}

	hc.cancel()
	hc.running = false
}

// AddTarget 添加监控目标
func (hc *HealthChecker) AddTarget(target string) error {
	u, err := url.Parse(target)
	if err != nil {
		return err
	}

	// 初始化为健康状态
	hc.statusMap.Store(u.String(), &backendStatus{
		URL:          u,
		Healthy:      true,
		FailCount:    0,
		SuccessCount: 0,
	})

	return nil
}

// RemoveTarget 移除监控目标
func (hc *HealthChecker) RemoveTarget(target string) error {
	u, err := url.Parse(target)
	if err != nil {
		return err
	}

	hc.statusMap.Delete(u.String())
	return nil
}

// IsHealthy 检查目标是否健康
func (hc *HealthChecker) IsHealthy(target string) bool {
	u, err := url.Parse(target)
	if err != nil {
		return false
	}

	value, ok := hc.statusMap.Load(u.String())
	if !ok {
		return false
	}

	status := value.(*backendStatus)
	return status.Healthy
}

// SetStatusChangeCallback 设置状态变更回调
func (hc *HealthChecker) SetStatusChangeCallback(callback func(string, bool)) {
	hc.statusChangeCallback = callback
}

// backendStatus 后端健康状态
type backendStatus struct {
	URL          *url.URL
	Healthy      bool
	FailCount    int
	SuccessCount int
}

// run 运行健康检查
func (hc *HealthChecker) run() {
	ticker := time.NewTicker(hc.config.Interval)
	defer ticker.Stop()

	for {
		select {
		case <-hc.ctx.Done():
			return
		case <-ticker.C:
			hc.checkAll()
		}
	}
}

// checkAll 检查所有后端
func (hc *HealthChecker) checkAll() {
	hc.statusMap.Range(func(key, value interface{}) bool {
		go hc.check(key.(string), value.(*backendStatus))
		return true
	})
}

// check 检查单个后端
func (hc *HealthChecker) check(key string, status *backendStatus) {
	// 创建检查请求
	u := *status.URL
	u.Path = hc.config.Path
	req, err := http.NewRequest(hc.config.Method, u.String(), nil)
	if err != nil {
		hc.updateStatus(key, status, false)
		return
	}

	// 设置超时的客户端
	client := &http.Client{
		Timeout: hc.config.Timeout,
		Transport: &http.Transport{
			DialContext: (&net.Dialer{
				Timeout:   hc.config.Timeout,
				KeepAlive: 30 * time.Second,
			}).DialContext,
			TLSHandshakeTimeout:   10 * time.Second,
			ResponseHeaderTimeout: hc.config.Timeout,
			ExpectContinueTimeout: 1 * time.Second,
			MaxIdleConns:          10,
			IdleConnTimeout:       30 * time.Second,
		},
	}

	// 发送请求
	resp, err := client.Do(req)
	if err != nil {
		hc.updateStatus(key, status, false)
		return
	}
	defer resp.Body.Close()

	// 检查响应状态
	if resp.StatusCode == hc.config.SuccessStatus {
		hc.updateStatus(key, status, true)
	} else {
		hc.updateStatus(key, status, false)
	}
}

// updateStatus 更新后端状态
func (hc *HealthChecker) updateStatus(key string, status *backendStatus, success bool) {
	if success {
		status.SuccessCount++
		status.FailCount = 0
		if !status.Healthy && status.SuccessCount >= hc.config.MinSuccess {
			status.Healthy = true
			if hc.statusChangeCallback != nil {
				hc.statusChangeCallback(key, true)
			}
		}
	} else {
		status.FailCount++
		status.SuccessCount = 0
		if status.Healthy && status.FailCount >= hc.config.MaxFails {
			status.Healthy = false
			if hc.statusChangeCallback != nil {
				hc.statusChangeCallback(key, false)
			}
		}
	}

	hc.statusMap.Store(key, status)
}