91 lines
2.7 KiB
Go
91 lines
2.7 KiB
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)
|
|||
|
|
}
|