xinmi/internal/application/service/wiki_provider_circuit.go

package service

import (
	"context"
	"fmt"
	"regexp"
	"strconv"
	"sync"
	"time"
)

const wikiProviderCooldown = 10 * time.Minute
const wikiProviderRateLimitCooldown = 60 * time.Second

var wikiPermanentStatus = regexp.MustCompile(`(?i)(?:status(?: code)?[ :=]+|http[ :=]+)(401|402|403|429)\b`)

type wikiProviderBlock struct {
	status int
	until  time.Time
}

// The local fallback also protects Lite mode and a temporarily unavailable Redis.
var wikiProviderBlocks sync.Map
var wikiProviderSlots sync.Map

// Bound a single process's Wiki fan-out across KBs sharing a tenant/model.
// Redis shares the cooldown across processes; these slots are process-local.
func acquireWikiProviderSlot(ctx context.Context, key string) (func(), error) {
	value, _ := wikiProviderSlots.LoadOrStore(key, make(chan struct{}, 2))
	slots := value.(chan struct{})
	select {
	case slots <- struct{}{}:
		return func() { <-slots }, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func wikiProviderCircuitKey(tenant uint64, model string) string {
	return fmt.Sprintf("wiki:provider-block:%d:%s", tenant, model)
}

func (s *wikiIngestService) checkWikiProvider(ctx context.Context, key string) error {
	if value, ok := wikiProviderBlocks.Load(key); ok {
		block := value.(wikiProviderBlock)
		if time.Now().Before(block.until) {
			return wikiProviderBlockedError(block.status)
		}
		wikiProviderBlocks.CompareAndDelete(key, value)
	}
	if s.redisClient != nil {
		status, err := s.redisClient.Get(ctx, key).Int()
		if err == nil && (status == 401 || status == 402 || status == 403 || status == 429) {
			return wikiProviderBlockedError(status)
		}
	}
	return nil
}

func (s *wikiIngestService) blockWikiProvider(ctx context.Context, key string, err error) {
	if err == nil {
		return
	}
	match := wikiPermanentStatus.FindStringSubmatch(err.Error())
	if len(match) != 2 {
		return
	}
	status, _ := strconv.Atoi(match[1])
	cooldown := wikiProviderCooldown
	if status == 429 {
		cooldown = wikiProviderRateLimitCooldown
	}
	block := wikiProviderBlock{status: status, until: time.Now().Add(cooldown)}
	wikiProviderBlocks.Store(key, block)
	time.AfterFunc(cooldown, func() { wikiProviderBlocks.CompareAndDelete(key, block) })
	if s.redisClient != nil {
		_ = s.redisClient.Set(ctx, key, status, cooldown).Err()
	}
}

func wikiProviderBlockedError(status int) error {
	if status == 429 {
		return fmt.Errorf("模型供应商限流（HTTP 429），Wiki 已暂停对此模型的自动请求 60 秒")
	}
	reason := "模型认证或访问权限异常"
	if status == 402 {
		reason = "模型供应商余额不足"
	}
	return fmt.Errorf("%s（HTTP %d），Wiki 已暂停对此模型的自动请求 10 分钟；请检查模型配置和供应商余额后重试", reason, status)
}