sub2api/backend/internal/service/account_traffic_policy.go

package service

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
)

const AccountTrafficPolicyKey = "account_traffic_control"

type AccountTrafficPolicy struct {
	StrictRPMEnabled     bool   `json:"strict_rpm_enabled"`
	RPM                  int    `json:"rpm"`
	Burst                int    `json:"burst"`
	AdaptiveEnabled      bool   `json:"adaptive_enabled"`
	AdaptiveMode         string `json:"adaptive_mode"`
	MinConcurrency       int    `json:"min_concurrency"`
	FailureThreshold     int    `json:"failure_threshold"`
	FailureWindowSeconds int    `json:"failure_window_seconds"`
	RecoverySeconds      int    `json:"recovery_seconds"`
}

func DefaultAccountTrafficPolicy() AccountTrafficPolicy {
	return AccountTrafficPolicy{RPM: 60, Burst: 5, AdaptiveMode: "observe", MinConcurrency: 1, FailureThreshold: 3, FailureWindowSeconds: 60, RecoverySeconds: 60}
}

func ParseAccountTrafficPolicy(extra map[string]any) (AccountTrafficPolicy, error) {
	p := DefaultAccountTrafficPolicy()
	value, exists := extra[AccountTrafficPolicyKey]
	if !exists || value == nil {
		return p, nil
	}
	encoded, err := json.Marshal(value)
	if err != nil {
		return p, err
	}
	d := json.NewDecoder(bytes.NewReader(encoded))
	d.DisallowUnknownFields()
	if err = d.Decode(&p); err != nil {
		return p, infraerrors.BadRequest("ACCOUNT_TRAFFIC_INVALID", "流量控制配置格式无效")
	}
	if p.RPM < 1 || p.RPM > 60000 || p.Burst < 1 || p.Burst > p.RPM || p.MinConcurrency < 1 || p.MinConcurrency > 10000 || p.FailureThreshold < 1 || p.FailureThreshold > 100 || p.FailureWindowSeconds < 10 || p.FailureWindowSeconds > 3600 || p.RecoverySeconds < 10 || p.RecoverySeconds > 3600 || (p.AdaptiveMode != "observe" && p.AdaptiveMode != "automatic") {
		return p, infraerrors.BadRequest("ACCOUNT_TRAFFIC_INVALID", "请检查 RPM、突发额度、并发下限、失败阈值与恢复周期；突发额度不能大于 RPM")
	}
	return p, nil
}

func (p AccountTrafficPolicy) Enabled() bool { return p.StrictRPMEnabled || p.AdaptiveEnabled }

func (p AccountTrafficPolicy) Enforces() bool {
	return p.StrictRPMEnabled || (p.AdaptiveEnabled && p.AdaptiveMode == "automatic")
}

type AccountTrafficPlan struct {
	AccountID int64
	Revision  int64
	Signature string
	HardLimit int
	Policy    AccountTrafficPolicy
}

func AccountTrafficPlanFor(a *Account) (AccountTrafficPlan, error) {
	if a == nil {
		return AccountTrafficPlan{}, nil
	}
	p, err := ParseAccountTrafficPolicy(a.Extra)
	if err != nil {
		return AccountTrafficPlan{}, err
	}
	limit := a.Mode1EffectiveConcurrency()
	if p.AdaptiveEnabled && limit <= 0 {
		return AccountTrafficPlan{}, infraerrors.BadRequest("ACCOUNT_TRAFFIC_INVALID", "启用自适应并发前，请先设置正数的账号并发上限")
	}
	encoded, _ := json.Marshal(struct {
		Policy AccountTrafficPolicy
		Limit  int
	}{p, limit})
	revision := int64(0)
	if !a.UpdatedAt.IsZero() {
		revision = a.UpdatedAt.UnixMicro()
	}
	return AccountTrafficPlan{AccountID: a.ID, Revision: revision, Signature: fmt.Sprintf("%x", sha256.Sum256(encoded)), HardLimit: limit, Policy: p}, nil
}

type accountTrafficPlanKey struct{}
type accountTrafficCoveredKey struct{}
type accountTrafficConfigErrorKey struct{}

// Attach only safe policy metadata, never account credentials. GET discovery
// and status polling are outside the model-request budget.
func WithAccountTrafficRequest(req *http.Request, a *Account) *http.Request {
	if req == nil || req.Method == http.MethodGet || req.Method == http.MethodHead {
		return req
	}
	if a == nil || a.Extra[AccountTrafficPolicyKey] == nil {
		return req
	}
	plan, err := AccountTrafficPlanFor(a)
	if err != nil {
		return req.WithContext(context.WithValue(req.Context(), accountTrafficConfigErrorKey{}, err))
	}
	if !plan.Policy.Enabled() {
		return req
	}
	return req.WithContext(context.WithValue(req.Context(), accountTrafficPlanKey{}, plan))
}

type AccountTrafficState struct {
	EffectiveConcurrency   int        `json:"effective_concurrency"`
	RecommendedConcurrency int        `json:"recommended_concurrency"`
	InFlight               int        `json:"in_flight"`
	RequestsLastMinute     int        `json:"requests_last_minute"`
	Accepted               int64      `json:"accepted"`
	RejectedRPM            int64      `json:"rejected_rpm"`
	RejectedConcurrency    int64      `json:"rejected_concurrency"`
	Upstream429            int64      `json:"upstream_429"`
	Upstream5xx            int64      `json:"upstream_5xx"`
	Completed              int64      `json:"completed"`
	AverageDurationMS      float64    `json:"average_duration_ms"`
	LastAdjustmentAt       *time.Time `json:"last_adjustment_at,omitempty"`
}

type AccountTrafficAdmission struct {
	Allowed         bool
	SkipObservation bool
	Reason          string
	RetryAfter      time.Duration
	State           AccountTrafficState
}

type AccountTrafficCache interface {
	Acquire(context.Context, AccountTrafficPlan, string) (AccountTrafficAdmission, error)
	Refresh(context.Context, int64, string) (bool, error)
	Finish(context.Context, AccountTrafficPlan, string, int, int64) error
	Snapshot(context.Context, AccountTrafficPlan) (AccountTrafficState, error)
	Sync(context.Context, AccountTrafficPlan) error
}

type AccountTrafficProvider interface{ AccountTrafficController() *AccountTrafficService }

type AccountTrafficLimitError struct {
	Reason     string
	RetryAfter time.Duration
	Status     int
}

func (e *AccountTrafficLimitError) Error() string { return "ACCOUNT_TRAFFIC_LIMIT: " + e.Reason }

func (e *AccountTrafficLimitError) FailoverError() *UpstreamFailoverError {
	status := e.Status
	if status == 0 {
		status = 429
	}
	seconds := int64((e.RetryAfter + time.Second - 1) / time.Second)
	if seconds < 1 {
		seconds = 1
	}
	errorType := AccountTrafficErrorType(status)
	body, _ := json.Marshal(map[string]any{"error": map[string]string{"type": errorType, "code": "account_traffic_limit", "message": e.Reason}})
	return &UpstreamFailoverError{StatusCode: status, ResponseBody: body, ResponseHeaders: http.Header{"Retry-After": []string{fmt.Sprint(seconds)}}, RequestScopedTransient: true, Scope: GatewayFailureScopeRequest, Reason: "account_traffic_limit", NextAccountAction: NextAccountStop, ClientStatusCode: status, ClientMessage: e.Reason}
}
func (e *AccountTrafficLimitError) As(target any) bool {
	if p, ok := target.(**UpstreamFailoverError); ok {
		*p = e.FailoverError()
		return true
	}
	return false
}

func AccountTrafficErrorType(status int) string {
	if status == http.StatusBadRequest {
		return "invalid_request_error"
	}
	if status >= 500 {
		return "api_error"
	}
	return "rate_limit_error"
}