junhong_cmp_fiber/internal/infrastructure/gatewaylimit/guard.go

package gatewaylimit

import (
	"context"
	"strconv"
	"time"

	"github.com/redis/go-redis/v9"
	"go.uber.org/zap"

	"github.com/break/junhong_cmp_fiber/internal/gateway"
	"github.com/break/junhong_cmp_fiber/pkg/constants"
)

// 分钟计数 Hash 的字段名。
const (
	statFieldWouldLimit  = "would_limit"
	statFieldDurationSum = "duration_ms_sum"
	statFieldAttempts    = "attempts"
	statPrefixExempt     = "exempt:"
	statPrefixUpstream   = "upstream:"
	statPrefixDuration   = "duration_le:"
)

// StatFieldRateLimitWait 是等待模式下进入限流等待的次数字段（等待超时另计入 rate_limited）。
const StatFieldRateLimitWait = "rate_wait"

// durationBuckets 是耗时分桶上界（毫秒），用于近似计算 p95。
var durationBuckets = []int64{100, 300, 500, 1000, 3000, 10000, 60000}

// Guard 实现 gateway.CallGuard：每次卡网关 HTTP 尝试前按运营商账户限流，结束后写分钟计数与卡最近错误。
type Guard struct {
	limiter    *Limiter
	resolver   *CarrierResolver
	classifier *Classifier
	redis      *redis.Client
	logger     *zap.Logger
}

// NewGuard 创建卡网关调用守卫。
func NewGuard(limiter *Limiter, resolver *CarrierResolver, classifier *Classifier, redisClient *redis.Client, logger *zap.Logger) *Guard {
	if logger == nil {
		logger = zap.NewNop()
	}
	return &Guard{limiter: limiter, resolver: resolver, classifier: classifier, redis: redisClient, logger: logger}
}

var _ gateway.CallGuard = (*Guard)(nil)

// BeforeAttempt 取限流名额；豁免调用直接放行。拦截模式下拿不到名额返回限流错误，本次方法调用随即结束。
func (g *Guard) BeforeAttempt(ctx context.Context, call gateway.CallInfo) error {
	if g == nil || g.limiter == nil || call.ExemptSource != "" {
		return nil
	}
	carrierID := g.carrierOf(ctx, call)
	decision, err := g.limiter.Check(ctx, carrierID, call.Operation)
	if err != nil {
		// 限流器故障时失败开放，避免 Redis 抖动阻断全部网关调用。
		g.logger.Warn("网关限流判断失败，本次放行", zap.Uint("carrier_id", carrierID), zap.String("operation", call.Operation), zap.Error(err))
		return nil
	}
	if decision.Exceed {
		g.incr(ctx, carrierID, call.Operation, statFieldWouldLimit)
	}
	if !decision.Allowed {
		// 等待模式（NSQ 消费者）：等名额直到超时；超时按限流错误返回，调用方按短延迟重排、不计失败。
		if maxWait := gateway.RateLimitWait(ctx); maxWait > 0 {
			g.incr(ctx, carrierID, call.Operation, StatFieldRateLimitWait)
			allowed, waitErr := g.limiter.Wait(ctx, carrierID, call.Operation, maxWait)
			if allowed {
				return nil
			}
			if waitErr != nil && ctx.Err() == nil {
				g.logger.Warn("网关限流等待失败，本次放行", zap.Uint("carrier_id", carrierID), zap.String("operation", call.Operation), zap.Error(waitErr))
				return nil
			}
		}
		g.incr(ctx, carrierID, call.Operation, constants.GatewayErrorRateLimited)
		return gateway.NewRateLimitedError()
	}
	return nil
}

// AfterAttempt 写分钟计数（按结果类别与耗时分桶）；查询失败时写卡该操作最近一次错误。
func (g *Guard) AfterAttempt(ctx context.Context, call gateway.CallInfo, callErr error, duration time.Duration) {
	if g == nil || g.redis == nil {
		return
	}
	ctx = context.WithoutCancel(ctx)
	carrierID := g.carrierOf(ctx, call)
	field := constants.GatewayResultSuccess
	var classification Classification
	gwErr, isGatewayErr := gateway.AsGatewayError(callErr)
	if callErr != nil {
		field = constants.GatewayErrorNetwork
		if isGatewayErr {
			field = gwErr.Category
			if gwErr.Category == constants.GatewayErrorUpstream {
				classification = g.classifier.Classify(ctx, call.Operation, gwErr.ProviderMessage)
				field = statPrefixUpstream + classification.Category
			}
		}
	}
	key := constants.RedisGatewayCallStatKey(carrierID, call.Operation, statMinute(time.Now()))
	durationMS := duration.Milliseconds()
	pipe := g.redis.Pipeline()
	pipe.HIncrBy(ctx, key, field, 1)
	pipe.HIncrBy(ctx, key, statFieldAttempts, 1)
	pipe.HIncrBy(ctx, key, statFieldDurationSum, durationMS)
	pipe.HIncrBy(ctx, key, statPrefixDuration+durationBucket(durationMS), 1)
	if call.ExemptSource != "" {
		pipe.HIncrBy(ctx, key, statPrefixExempt+call.ExemptSource, 1)
	}
	pipe.Expire(ctx, key, constants.GatewayCallStatTTL)
	if callErr != nil && isQueryOperation(call.Operation) && call.CardNo != "" {
		errKey := constants.RedisGatewayCardLastErrorKey(call.CardNo, call.Operation)
		values := map[string]any{"category": field, "at": time.Now().Format(time.RFC3339)}
		if isGatewayErr {
			values["provider_code"] = strconv.Itoa(gwErr.ProviderCode)
			values["provider_message"] = gwErr.ProviderMessage
		}
		pipe.HSet(ctx, errKey, values)
		pipe.Expire(ctx, errKey, constants.GatewayCardLastErrorTTL)
	}
	if _, err := pipe.Exec(ctx); err != nil {
		g.logger.Debug("写网关调用统计失败", zap.Error(err))
	}
}

// Classify 暴露分类器供停复机记录等调用方对上游错误分类。
func (g *Guard) Classify(ctx context.Context, op, message string) Classification {
	if g == nil {
		return Classification{Category: constants.GatewayClassBusiness}
	}
	return g.classifier.Classify(ctx, op, message)
}

func (g *Guard) carrierOf(ctx context.Context, call gateway.CallInfo) uint {
	if call.CarrierID != 0 || g.resolver == nil {
		return call.CarrierID
	}
	carrierID, err := g.resolver.Resolve(ctx, call.CardNo)
	if err != nil {
		g.logger.Debug("解析卡所属运营商账户失败，按未知账户计量", zap.String("iccid", call.CardNo), zap.Error(err))
	}
	return carrierID
}

func (g *Guard) incr(ctx context.Context, carrierID uint, op, field string) {
	if g.redis == nil {
		return
	}
	key := constants.RedisGatewayCallStatKey(carrierID, op, statMinute(time.Now()))
	pipe := g.redis.Pipeline()
	pipe.HIncrBy(ctx, key, field, 1)
	pipe.Expire(ctx, key, constants.GatewayCallStatTTL)
	_, _ = pipe.Exec(context.WithoutCancel(ctx))
}

func isQueryOperation(op string) bool {
	return op == constants.GatewayOpQueryFlow || op == constants.GatewayOpQueryStatus || op == constants.GatewayOpQueryRealname
}

// statMinute 返回上海时区的分钟标识（yyyyMMddHHmm）。
func statMinute(t time.Time) string {
	return t.In(shanghai()).Format("200601021504")
}

func durationBucket(ms int64) string {
	for _, upper := range durationBuckets {
		if ms <= upper {
			return strconv.FormatInt(upper, 10)
		}
	}
	return "inf"
}

func shanghai() *time.Location {
	location, err := time.LoadLocation("Asia/Shanghai")
	if err != nil {
		return time.FixedZone("CST", 8*3600)
	}
	return location
}