Files

91 lines
2.7 KiB
Go
Raw Permalink Normal View History

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