frist

internal/selector/selector.go

@@ -0,0 +1,177 @@
+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)
+}