kami_gateway/internal/proxy/proxy_pool.go

package proxy

import (
	"context"
	"fmt"
	"gateway/internal/otelTrace"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/duke-git/lancet/v2/slice"
	"go.uber.org/zap"
)

// Proxy 表示一个代理服务器
type Proxy struct {
	URL         string
	LastTest    time.Time
	IsValid     bool
	UseCount    int64     // 使用次数
	LastUseTime time.Time // 最后使用时间
	TotalTime   int64     // 总使用时间(毫秒)
}

// Pool 代理池
type Pool struct {
	proxies []*Proxy
	mu      sync.RWMutex
	config  *Config
}

var (
	globalProxyPool *Pool
)

// GetGlobalProxyPool 获取全局代理池实例
func GetGlobalProxyPool() *Pool {
	sync.OnceFunc(func() {
		globalProxyPool = NewProxyPool()
	})()
	return globalProxyPool
}

// NewProxyPool 创建一个新的代理池
func NewProxyPool() *Pool {
	return &Pool{
		proxies: make([]*Proxy, 0),
		config:  DefaultConfig(),
	}
}

// InitWithConfig 使用配置初始化代理池
func (p *Pool) InitWithConfig(config *Config) {
	p.mu.Lock()
	defer p.mu.Unlock()

	p.config = config
	p.proxies = make([]*Proxy, 0, len(config.Proxies))

	for _, proxyURL := range config.Proxies {
		p.proxies = append(p.proxies, &Proxy{
			URL:         proxyURL,
			LastTest:    time.Now(),
			IsValid:     true,
			UseCount:    0,
			LastUseTime: time.Time{},
			TotalTime:   0,
		})
	}
}

// AddProxy 添加代理到代理池
func (p *Pool) AddProxy(proxyURL string) {
	p.mu.Lock()
	defer p.mu.Unlock()

	p.proxies = append(p.proxies, &Proxy{
		URL:         proxyURL,
		LastTest:    time.Now(),
		IsValid:     true,
		UseCount:    0,
		LastUseTime: time.Time{},
		TotalTime:   0,
	})
}

// GetRandomProxy 随机获取一个可用的代理
func (p *Pool) GetRandomProxy() (*Proxy, error) {
	p.mu.RLock()
	defer p.mu.RUnlock()

	if len(p.proxies) == 0 {
		return nil, fmt.Errorf("no proxies available")
	}

	slice.SortBy(p.proxies, func(a, b *Proxy) bool {
		return a.UseCount < b.UseCount
	})

	var selectedProxy *Proxy

	for _, proxy := range p.proxies {
		err := p.TestProxy(proxy)
		if err != nil {
			otelTrace.Logger.WithContext(context.Background()).Info("proxy失效", zap.Any("proxy", proxy))
			proxy.IsValid = false
			continue
		}
		if proxy.IsValid && time.Since(proxy.LastTest) < time.Duration(p.config.Timeout)*time.Second {
			selectedProxy = proxy
			break
		}
	}

	return selectedProxy, nil
}

// MarkProxyUsed 标记代理被使用
func (p *Pool) MarkProxyUsed(proxy *Proxy, duration time.Duration) {
	p.mu.Lock()
	defer p.mu.Unlock()

	proxy.UseCount++
	proxy.LastUseTime = time.Now()
	proxy.TotalTime += duration.Milliseconds()
}

// TestProxy 测试代理是否可用
func (p *Pool) TestProxy(proxy *Proxy) error {
	proxyURL, err := url.Parse(proxy.URL)
	if err != nil {
		proxy.LastTest = time.Now()
		proxy.IsValid = false
		return err
	}

	client := &http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyURL(proxyURL),
		},
		Timeout: time.Duration(p.config.Timeout) * time.Second,
	}

	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(p.config.Timeout)*time.Second)
	defer cancel()

	req, err := http.NewRequestWithContext(ctx, "GET", p.config.TestURL, nil)
	if err != nil {
		proxy.LastTest = time.Now()
		proxy.IsValid = false
		return err
	}

	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		proxy.LastTest = time.Now()
		proxy.IsValid = false
		return fmt.Errorf("proxy test failed with status code: %d", resp.StatusCode)
	}

	proxy.LastTest = time.Now()
	proxy.IsValid = true
	return nil
}

// UpdateProxies 更新代理池中所有代理的状态
func (p *Pool) UpdateProxies() {
	p.mu.Lock()
	defer p.mu.Unlock()

	for _, proxy := range p.proxies {
		if time.Since(proxy.LastTest) >= time.Duration(p.config.Timeout)*time.Second {
			err := p.TestProxy(proxy)
			if err != nil {
				proxy.IsValid = false
			}
		}
	}
}

// GetProxyStats 获取代理统计信息
func (p *Pool) GetProxyStats() []map[string]any {
	p.mu.RLock()
	defer p.mu.RUnlock()

	stats := make([]map[string]interface{}, 0, len(p.proxies))
	for _, proxy := range p.proxies {
		stat := map[string]interface{}{
			"url":           proxy.URL,
			"is_valid":      proxy.IsValid,
			"use_count":     proxy.UseCount,
			"last_use_time": proxy.LastUseTime,
			"total_time_ms": proxy.TotalTime,
			"last_test":     proxy.LastTest,
		}
		stats = append(stats, stat)
	}
	return stats
}