178 lines
3.4 KiB
Go
178 lines
3.4 KiB
Go
package selector
|
||
|
||
import (
|
||
"context"
|
||
"crypto/rand"
|
||
"encoding/hex"
|
||
"math/big"
|
||
mathrand "math/rand"
|
||
"net/url"
|
||
"strings"
|
||
"time"
|
||
|
||
"proxyrotator/internal/model"
|
||
"proxyrotator/internal/store"
|
||
)
|
||
|
||
// Selector 代理选择器
|
||
type Selector struct {
|
||
store store.ProxyStore
|
||
leaseTTL time.Duration
|
||
}
|
||
|
||
// NewSelector 创建选择器
|
||
func NewSelector(store store.ProxyStore, leaseTTL time.Duration) *Selector {
|
||
if leaseTTL <= 0 {
|
||
leaseTTL = 60 * time.Second
|
||
}
|
||
return &Selector{
|
||
store: store,
|
||
leaseTTL: leaseTTL,
|
||
}
|
||
}
|
||
|
||
// Next 获取下一个可用代理
|
||
func (s *Selector) Next(ctx context.Context, req model.SelectRequest) (*model.Lease, error) {
|
||
policy := req.Policy
|
||
if policy == "" {
|
||
policy = "round_robin"
|
||
}
|
||
|
||
// 查询可用代理
|
||
proxies, err := s.store.List(ctx, model.ProxyQuery{
|
||
Group: req.Group,
|
||
TagsAny: req.TagsAny,
|
||
StatusIn: []model.ProxyStatus{model.StatusAlive},
|
||
OnlyEnabled: true,
|
||
Limit: 5000,
|
||
})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
if len(proxies) == 0 {
|
||
return nil, model.ErrNoProxy
|
||
}
|
||
|
||
// 根据策略选择
|
||
var chosen model.Proxy
|
||
switch policy {
|
||
case "round_robin":
|
||
key := "rr:" + req.Group + ":" + normalizeSite(req.Site)
|
||
idx, err := s.store.NextIndex(ctx, key, len(proxies))
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
chosen = proxies[idx]
|
||
|
||
case "random":
|
||
idx := mathrand.Intn(len(proxies))
|
||
chosen = proxies[idx]
|
||
|
||
case "weighted":
|
||
chosen = weightedPickByScore(proxies)
|
||
|
||
default:
|
||
return nil, model.ErrBadPolicy
|
||
}
|
||
|
||
// 创建租约
|
||
lease := model.Lease{
|
||
LeaseID: newLeaseID(),
|
||
ProxyID: chosen.ID,
|
||
Proxy: chosen,
|
||
Group: req.Group,
|
||
Site: req.Site,
|
||
ExpireAt: time.Now().Add(s.leaseTTL),
|
||
}
|
||
|
||
// 尝试保存租约（失败也可降级不存）
|
||
_ = s.store.CreateLease(ctx, lease)
|
||
|
||
return &lease, nil
|
||
}
|
||
|
||
// Report 上报使用结果
|
||
func (s *Selector) Report(ctx context.Context, leaseID, proxyID string, success bool, latencyMs int64, errText string) error {
|
||
now := time.Now()
|
||
|
||
if success {
|
||
status := model.StatusAlive
|
||
return s.store.UpdateHealth(ctx, proxyID, model.HealthPatch{
|
||
Status: &status,
|
||
ScoreDelta: 1,
|
||
SuccessInc: 1,
|
||
LatencyMs: &latencyMs,
|
||
CheckedAt: &now,
|
||
})
|
||
}
|
||
|
||
status := model.StatusDead
|
||
return s.store.UpdateHealth(ctx, proxyID, model.HealthPatch{
|
||
Status: &status,
|
||
ScoreDelta: -3,
|
||
FailInc: 1,
|
||
CheckedAt: &now,
|
||
})
|
||
}
|
||
|
||
// normalizeSite 规范化站点 URL（提取域名）
|
||
func normalizeSite(site string) string {
|
||
if site == "" {
|
||
return "default"
|
||
}
|
||
|
||
u, err := url.Parse(site)
|
||
if err != nil {
|
||
return site
|
||
}
|
||
|
||
host := u.Hostname()
|
||
if host == "" {
|
||
return site
|
||
}
|
||
|
||
// 去除 www 前缀
|
||
host = strings.TrimPrefix(host, "www.")
|
||
return host
|
||
}
|
||
|
||
// weightedPickByScore 按分数加权随机选择
|
||
func weightedPickByScore(proxies []model.Proxy) model.Proxy {
|
||
// 计算权重（分数 + 偏移量确保正数）
|
||
const offset = 100
|
||
totalWeight := 0
|
||
weights := make([]int, len(proxies))
|
||
|
||
for i, p := range proxies {
|
||
w := p.Score + offset
|
||
if w < 1 {
|
||
w = 1
|
||
}
|
||
weights[i] = w
|
||
totalWeight += w
|
||
}
|
||
|
||
// 随机选择
|
||
r := mathrand.Intn(totalWeight)
|
||
for i, w := range weights {
|
||
r -= w
|
||
if r < 0 {
|
||
return proxies[i]
|
||
}
|
||
}
|
||
|
||
return proxies[0]
|
||
}
|
||
|
||
// newLeaseID 生成租约 ID
|
||
func newLeaseID() string {
|
||
b := make([]byte, 16)
|
||
if _, err := rand.Read(b); err != nil {
|
||
// fallback
|
||
n, _ := rand.Int(rand.Reader, big.NewInt(1<<62))
|
||
return "lease_" + n.String()
|
||
}
|
||
return "lease_" + hex.EncodeToString(b)
|
||
}
|