Blizzard
529cd0fdd6
根治 core NATS 不回放导致的"连晚/刷新/抖动重连丢实时 token"： - store/redis：token 流落 Redis Stream（XADD，带 TTL）——StreamAppend / StreamRead(XREAD BLOCK，支持从任意 ID 续读)。 - gateway：SubmitTask 起即启动 token 录制器（订阅早于 dispatcher 产 token， 全量捕获，与 SSE 客户端是否在线无关）；StreamTask 改从 Redis Stream 读， 发 SSE 时带 id（= Redis entry ID），浏览器重连自动带 Last-Event-ID 续传。 Redis 降级时回退原 live-NATS 路径（streamTaskLive）。 实测：任务跑完后才连 SSE → 完整回放；带 Last-Event-ID 重连 → 断点续传不重不漏； make test-go 全绿。最终答复本就持久在 Redis 历史，此改让"实时流"也无损。 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
115 lines
3.3 KiB
Go
115 lines
3.3 KiB
Go
package store
|
||
|
||
import (
|
||
"context"
|
||
"log"
|
||
"time"
|
||
|
||
"github.com/redis/go-redis/v9"
|
||
)
|
||
|
||
// Redis 持有 CacheDB 连接（Session / Rate Limit）。
|
||
// rdb 为 nil 表示降级模式（连接失败时放行，不限流）。
|
||
type Redis struct {
|
||
rdb *redis.Client
|
||
}
|
||
|
||
// OpenRedis 连接 CacheDB。连接失败不 fatal：返回降级实例（限流放行）。
|
||
func OpenRedis(addr string) *Redis {
|
||
rdb := redis.NewClient(&redis.Options{Addr: addr})
|
||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||
defer cancel()
|
||
if err := rdb.Ping(ctx).Err(); err != nil {
|
||
log.Printf("[store] redis 不可用，降级运行（不限流）: %v", err)
|
||
_ = rdb.Close()
|
||
return &Redis{}
|
||
}
|
||
log.Println("[store] redis connected")
|
||
return &Redis{rdb: rdb}
|
||
}
|
||
|
||
// Enabled 报告是否处于真实限流模式。
|
||
func (r *Redis) Enabled() bool { return r.rdb != nil }
|
||
|
||
// Allow 滑动窗口计数限流：在 window 内对 key 累加，超过 limit 即拒绝。
|
||
// 降级模式（rdb==nil）始终放行。
|
||
func (r *Redis) Allow(ctx context.Context, key string, limit int64, window time.Duration) (bool, error) {
|
||
if r.rdb == nil {
|
||
return true, nil
|
||
}
|
||
rk := "sundynix:ratelimit:" + key
|
||
n, err := r.rdb.Incr(ctx, rk).Result()
|
||
if err != nil {
|
||
return true, err // 限流后端故障时放行，不阻断业务
|
||
}
|
||
if n == 1 {
|
||
_ = r.rdb.Expire(ctx, rk, window).Err() // 首次计数设置窗口过期
|
||
}
|
||
return n <= limit, nil
|
||
}
|
||
|
||
// ---- Token 流持久化（Redis Stream：可回放的追加日志，根治 SSE 连晚/重连丢 token）----
|
||
|
||
const streamTTL = 10 * time.Minute
|
||
|
||
func streamKey(taskID string) string { return "sundynix:stream:" + taskID }
|
||
|
||
// StreamEntry 是 token 流里的一条记录（ID 用于 SSE 的 Last-Event-ID 断点续传）。
|
||
type StreamEntry struct {
|
||
ID string
|
||
Kind string // token / done
|
||
Data string
|
||
}
|
||
|
||
// StreamAppend 把一条 token / 结束标记追加到任务的 Redis Stream（带 TTL 自动清理）。
|
||
func (r *Redis) StreamAppend(ctx context.Context, taskID, kind, data string) error {
|
||
if r.rdb == nil {
|
||
return nil
|
||
}
|
||
k := streamKey(taskID)
|
||
if err := r.rdb.XAdd(ctx, &redis.XAddArgs{Stream: k, Values: map[string]any{"kind": kind, "data": data}}).Err(); err != nil {
|
||
return err
|
||
}
|
||
return r.rdb.Expire(ctx, k, streamTTL).Err()
|
||
}
|
||
|
||
// StreamRead 从 lastID 之后阻塞读取新条目（XREAD BLOCK）。lastID="0" 表示从头回放。
|
||
// 阻塞超时无新数据时返回空切片 + 原 lastID（调用方据此继续轮询）。
|
||
func (r *Redis) StreamRead(ctx context.Context, taskID, lastID string, block time.Duration) ([]StreamEntry, string, error) {
|
||
if r.rdb == nil {
|
||
return nil, lastID, nil
|
||
}
|
||
res, err := r.rdb.XRead(ctx, &redis.XReadArgs{
|
||
Streams: []string{streamKey(taskID), lastID}, Block: block, Count: 256,
|
||
}).Result()
|
||
if err == redis.Nil {
|
||
return nil, lastID, nil // 阻塞超时、无新条目
|
||
}
|
||
if err != nil {
|
||
return nil, lastID, err
|
||
}
|
||
var out []StreamEntry
|
||
nl := lastID
|
||
for _, st := range res {
|
||
for _, m := range st.Messages {
|
||
out = append(out, StreamEntry{ID: m.ID, Kind: asString(m.Values["kind"]), Data: asString(m.Values["data"])})
|
||
nl = m.ID
|
||
}
|
||
}
|
||
return out, nl, nil
|
||
}
|
||
|
||
func asString(v any) string {
|
||
if s, ok := v.(string); ok {
|
||
return s
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// Close 释放底层连接。
|
||
func (r *Redis) Close() {
|
||
if r.rdb != nil {
|
||
_ = r.rdb.Close()
|
||
}
|
||
}
|