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) }