chore: refresh gramsrv public release
This commit is contained in:
A
parent
75cebe8dbf
commit
70b6820474
1274 changed files with 378751 additions and 59919 deletions
|
|
@ -10,21 +10,11 @@ import (
|
|||
"telesrv/internal/store"
|
||||
)
|
||||
|
||||
// PtsAllocator 用 Redis INCR 分配账号级 pts，Redis 丢失时从 PG durable log 恢复当前最大 pts。
|
||||
type PtsAllocator struct {
|
||||
counter counterAllocator
|
||||
}
|
||||
|
||||
// BoxIDAllocator 用 Redis INCR 分配 owner 视角的 message box id。
|
||||
type BoxIDAllocator struct {
|
||||
counter counterAllocator
|
||||
}
|
||||
|
||||
// ChannelPtsAllocator 用 Redis INCR 分配 channel 维度 pts。
|
||||
type ChannelPtsAllocator struct {
|
||||
counter counterAllocator
|
||||
}
|
||||
|
||||
// ChannelIDAllocator 用 Redis INCR 分配全局 channel/supergroup id。
|
||||
type ChannelIDAllocator struct {
|
||||
counter counterAllocator
|
||||
|
|
@ -35,6 +25,11 @@ type ChannelMessageIDAllocator struct {
|
|||
counter counterAllocator
|
||||
}
|
||||
|
||||
// SecretChatIDAllocator 用 Redis INCR 分配全局 secret chat id（int32 量级）。
|
||||
type SecretChatIDAllocator struct {
|
||||
counter counterAllocator
|
||||
}
|
||||
|
||||
type counterAllocator struct {
|
||||
c *redis.Client
|
||||
source store.CounterSource
|
||||
|
|
@ -51,14 +46,6 @@ if current then
|
|||
return redis.call("INCR", KEYS[1])
|
||||
end
|
||||
return -1
|
||||
`)
|
||||
|
||||
counterNextByScript = redis.NewScript(`
|
||||
local current = redis.call("GET", KEYS[1])
|
||||
if current then
|
||||
return redis.call("INCRBY", KEYS[1], ARGV[1])
|
||||
end
|
||||
return -1
|
||||
`)
|
||||
|
||||
counterRecoverCurrentScript = redis.NewScript(`
|
||||
|
|
@ -78,25 +65,24 @@ end
|
|||
return redis.call("INCR", KEYS[1])
|
||||
`)
|
||||
|
||||
counterRecoverNextByScript = redis.NewScript(`
|
||||
counterNextAtLeastScript = redis.NewScript(`
|
||||
local current = redis.call("GET", KEYS[1])
|
||||
if not current then
|
||||
if (not current) or tonumber(current) < tonumber(ARGV[1]) then
|
||||
redis.call("SET", KEYS[1], ARGV[1])
|
||||
end
|
||||
return redis.call("INCRBY", KEYS[1], ARGV[2])
|
||||
return redis.call("INCR", KEYS[1])
|
||||
`)
|
||||
|
||||
counterSetAtLeastScript = redis.NewScript(`
|
||||
local current = redis.call("GET", KEYS[1])
|
||||
if (not current) or tonumber(current) < tonumber(ARGV[1]) then
|
||||
redis.call("SET", KEYS[1], ARGV[1])
|
||||
return tonumber(ARGV[1])
|
||||
end
|
||||
return tonumber(current)
|
||||
`)
|
||||
)
|
||||
|
||||
// NewPtsAllocator 创建 Redis-backed pts allocator。
|
||||
func NewPtsAllocator(c *redis.Client, source store.CounterSource) *PtsAllocator {
|
||||
return &PtsAllocator{counter: counterAllocator{
|
||||
c: c,
|
||||
source: source,
|
||||
key: ptsKey,
|
||||
name: "pts",
|
||||
}}
|
||||
}
|
||||
|
||||
// NewBoxIDAllocator 创建 Redis-backed message box id allocator。
|
||||
func NewBoxIDAllocator(c *redis.Client, source store.CounterSource) *BoxIDAllocator {
|
||||
return &BoxIDAllocator{counter: counterAllocator{
|
||||
|
|
@ -107,16 +93,6 @@ func NewBoxIDAllocator(c *redis.Client, source store.CounterSource) *BoxIDAlloca
|
|||
}}
|
||||
}
|
||||
|
||||
// NewChannelPtsAllocator 创建 Redis-backed channel pts allocator。
|
||||
func NewChannelPtsAllocator(c *redis.Client, source store.CounterSource) *ChannelPtsAllocator {
|
||||
return &ChannelPtsAllocator{counter: counterAllocator{
|
||||
c: c,
|
||||
source: source,
|
||||
key: channelPtsKey,
|
||||
name: "channel_pts",
|
||||
}}
|
||||
}
|
||||
|
||||
// NewChannelIDAllocator 创建 Redis-backed channel id allocator。
|
||||
func NewChannelIDAllocator(c *redis.Client, source store.CounterSource) *ChannelIDAllocator {
|
||||
return &ChannelIDAllocator{counter: counterAllocator{
|
||||
|
|
@ -137,41 +113,32 @@ func NewChannelMessageIDAllocator(c *redis.Client, source store.CounterSource) *
|
|||
}}
|
||||
}
|
||||
|
||||
func ptsKey(userID int64) string {
|
||||
return fmt.Sprintf("counter:pts:{%d}", userID)
|
||||
// NewSecretChatIDAllocator 创建 Redis-backed secret chat id allocator。
|
||||
func NewSecretChatIDAllocator(c *redis.Client, source store.CounterSource) *SecretChatIDAllocator {
|
||||
return &SecretChatIDAllocator{counter: counterAllocator{
|
||||
c: c,
|
||||
source: source,
|
||||
key: secretChatIDKey,
|
||||
name: "secret_chat_id",
|
||||
}}
|
||||
}
|
||||
|
||||
func boxIDKey(userID int64) string {
|
||||
return fmt.Sprintf("counter:box_id:{%d}", userID)
|
||||
}
|
||||
|
||||
func channelPtsKey(channelID int64) string {
|
||||
return fmt.Sprintf("counter:channel_pts:{%d}", channelID)
|
||||
}
|
||||
|
||||
func channelIDKey(_ int64) string {
|
||||
return "counter:channel_id"
|
||||
}
|
||||
|
||||
func secretChatIDKey(_ int64) string {
|
||||
return "counter:secret_chat_id"
|
||||
}
|
||||
|
||||
func channelMessageIDKey(channelID int64) string {
|
||||
return fmt.Sprintf("counter:channel_msg_id:{%d}", channelID)
|
||||
}
|
||||
|
||||
func (a *PtsAllocator) NextPts(ctx context.Context, userID int64) (int, error) {
|
||||
v, err := a.counter.next(ctx, userID)
|
||||
return int(v), err
|
||||
}
|
||||
|
||||
func (a *PtsAllocator) NextPtsN(ctx context.Context, userID int64, count int) (int, error) {
|
||||
v, err := a.counter.nextBy(ctx, userID, count)
|
||||
return int(v), err
|
||||
}
|
||||
|
||||
func (a *PtsAllocator) CurrentPts(ctx context.Context, userID int64) (int, error) {
|
||||
v, err := a.counter.current(ctx, userID)
|
||||
return int(v), err
|
||||
}
|
||||
|
||||
func (a *BoxIDAllocator) NextBoxID(ctx context.Context, userID int64) (int, error) {
|
||||
v, err := a.counter.next(ctx, userID)
|
||||
return int(v), err
|
||||
|
|
@ -182,29 +149,64 @@ func (a *BoxIDAllocator) CurrentBoxID(ctx context.Context, userID int64) (int, e
|
|||
return int(v), err
|
||||
}
|
||||
|
||||
func (a *ChannelPtsAllocator) NextChannelPts(ctx context.Context, channelID int64) (int, error) {
|
||||
v, err := a.counter.next(ctx, channelID)
|
||||
return int(v), err
|
||||
}
|
||||
|
||||
func (a *ChannelPtsAllocator) NextChannelPtsN(ctx context.Context, channelID int64, count int) (int, error) {
|
||||
v, err := a.counter.nextBy(ctx, channelID, count)
|
||||
return int(v), err
|
||||
}
|
||||
|
||||
func (a *ChannelPtsAllocator) CurrentChannelPts(ctx context.Context, channelID int64) (int, error) {
|
||||
v, err := a.counter.current(ctx, channelID)
|
||||
return int(v), err
|
||||
// BumpBoxIDAtLeast advances the Redis box id counter to at least floor without
|
||||
// allocating a visible id. It is a cold-path self-heal for Redis counters that
|
||||
// lag behind message_boxes after external/dev writes.
|
||||
func (a *BoxIDAllocator) BumpBoxIDAtLeast(ctx context.Context, userID int64, floor int) error {
|
||||
if a.counter.c == nil {
|
||||
return fmt.Errorf("redis box_id counter: nil client")
|
||||
}
|
||||
if _, err := counterSetAtLeastScript.Run(ctx, a.counter.c, []string{boxIDKey(userID)}, floor).Int64(); err != nil {
|
||||
return fmt.Errorf("redis set-at-least box_id counter: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *ChannelIDAllocator) NextChannelID(ctx context.Context) (int64, error) {
|
||||
return a.counter.next(ctx, 1)
|
||||
}
|
||||
|
||||
// NextChannelIDAtLeast 把计数器至少顶到 floor 后再分配下一个 id。
|
||||
// 用于撞主键自愈：Redis 快照回退或测试 fallback 分配器绕过 Redis 写库
|
||||
// 后，计数器可能落后于 channels 表真实最大 id。
|
||||
func (a *ChannelIDAllocator) NextChannelIDAtLeast(ctx context.Context, floor int64) (int64, error) {
|
||||
if a.counter.c == nil {
|
||||
return 0, fmt.Errorf("redis channel_id counter: nil client")
|
||||
}
|
||||
v, err := counterNextAtLeastScript.Run(ctx, a.counter.c, []string{channelIDKey(1)}, floor).Int64()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("redis next-at-least channel_id counter: %w", err)
|
||||
}
|
||||
return v, nil
|
||||
}
|
||||
|
||||
func (a *ChannelIDAllocator) CurrentChannelID(ctx context.Context) (int64, error) {
|
||||
return a.counter.current(ctx, 1)
|
||||
}
|
||||
|
||||
func (a *SecretChatIDAllocator) NextSecretChatID(ctx context.Context) (int, error) {
|
||||
v, err := a.counter.next(ctx, 1)
|
||||
return int(v), err
|
||||
}
|
||||
|
||||
// NextSecretChatIDAtLeast 把计数器至少顶到 floor 后再分配下一个 id（撞 chat_id
|
||||
// 主键自愈：Redis 快照回退或外部写库后计数器落后于 secret_chats 表最大 id）。
|
||||
func (a *SecretChatIDAllocator) NextSecretChatIDAtLeast(ctx context.Context, floor int) (int, error) {
|
||||
if a.counter.c == nil {
|
||||
return 0, fmt.Errorf("redis secret_chat_id counter: nil client")
|
||||
}
|
||||
v, err := counterNextAtLeastScript.Run(ctx, a.counter.c, []string{secretChatIDKey(1)}, floor).Int64()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("redis next-at-least secret_chat_id counter: %w", err)
|
||||
}
|
||||
return int(v), nil
|
||||
}
|
||||
|
||||
func (a *SecretChatIDAllocator) CurrentSecretChatID(ctx context.Context) (int, error) {
|
||||
v, err := a.counter.current(ctx, 1)
|
||||
return int(v), err
|
||||
}
|
||||
|
||||
func (a *ChannelMessageIDAllocator) NextChannelMessageID(ctx context.Context, channelID int64) (int, error) {
|
||||
v, err := a.counter.next(ctx, channelID)
|
||||
return int(v), err
|
||||
|
|
@ -216,24 +218,11 @@ func (a *ChannelMessageIDAllocator) CurrentChannelMessageID(ctx context.Context,
|
|||
}
|
||||
|
||||
func (a counterAllocator) next(ctx context.Context, userID int64) (int64, error) {
|
||||
return a.nextBy(ctx, userID, 1)
|
||||
}
|
||||
|
||||
func (a counterAllocator) nextBy(ctx context.Context, userID int64, count int) (int64, error) {
|
||||
if count <= 0 {
|
||||
return 0, fmt.Errorf("redis next %s counter: invalid count %d", a.name, count)
|
||||
}
|
||||
key, err := a.validatedKey(userID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
script := counterNextScript
|
||||
args := []any{}
|
||||
if count > 1 {
|
||||
script = counterNextByScript
|
||||
args = append(args, count)
|
||||
}
|
||||
v, err := script.Run(ctx, a.c, []string{key}, args...).Int64()
|
||||
v, err := counterNextScript.Run(ctx, a.c, []string{key}).Int64()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("redis next %s counter: %w", a.name, err)
|
||||
}
|
||||
|
|
@ -244,13 +233,7 @@ func (a counterAllocator) nextBy(ctx context.Context, userID int64, count int) (
|
|||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
recoverScript := counterRecoverNextScript
|
||||
recoverArgs := []any{recovered}
|
||||
if count > 1 {
|
||||
recoverScript = counterRecoverNextByScript
|
||||
recoverArgs = append(recoverArgs, count)
|
||||
}
|
||||
v, err = recoverScript.Run(ctx, a.c, []string{key}, recoverArgs...).Int64()
|
||||
v, err = counterRecoverNextScript.Run(ctx, a.c, []string{key}, recovered).Int64()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("redis recover-next %s counter: %w", a.name, err)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ func (s staticCounterSource) Current(context.Context, int64) (int, error) {
|
|||
return s.value, nil
|
||||
}
|
||||
|
||||
func TestRedisAllocatorsRecoverFromCounterSource(t *testing.T) {
|
||||
func TestRedisBoxAllocatorRecoverFromCounterSource(t *testing.T) {
|
||||
addr := os.Getenv("TELESRV_TEST_REDIS_ADDR")
|
||||
if addr == "" {
|
||||
t.Skip("set TELESRV_TEST_REDIS_ADDR to run redis integration test")
|
||||
|
|
@ -29,23 +29,7 @@ func TestRedisAllocatorsRecoverFromCounterSource(t *testing.T) {
|
|||
t.Cleanup(func() { _ = c.Close() })
|
||||
|
||||
userID := time.Now().UnixNano()
|
||||
t.Cleanup(func() { _ = c.Del(ctx, ptsKey(userID), boxIDKey(userID)).Err() })
|
||||
|
||||
pts := NewPtsAllocator(c, staticCounterSource{value: 41})
|
||||
currentPts, err := pts.CurrentPts(ctx, userID)
|
||||
if err != nil {
|
||||
t.Fatalf("CurrentPts: %v", err)
|
||||
}
|
||||
if currentPts != 41 {
|
||||
t.Fatalf("current pts = %d, want recovered 41", currentPts)
|
||||
}
|
||||
nextPts, err := pts.NextPts(ctx, userID)
|
||||
if err != nil {
|
||||
t.Fatalf("NextPts: %v", err)
|
||||
}
|
||||
if nextPts != 42 {
|
||||
t.Fatalf("next pts = %d, want 42", nextPts)
|
||||
}
|
||||
t.Cleanup(func() { _ = c.Del(ctx, boxIDKey(userID)).Err() })
|
||||
|
||||
boxes := NewBoxIDAllocator(c, staticCounterSource{value: 100})
|
||||
currentBox, err := boxes.CurrentBoxID(ctx, userID)
|
||||
|
|
@ -64,7 +48,7 @@ func TestRedisAllocatorsRecoverFromCounterSource(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestRedisAllocatorConcurrentFirstUse(t *testing.T) {
|
||||
func TestRedisBoxAllocatorConcurrentFirstUse(t *testing.T) {
|
||||
addr := os.Getenv("TELESRV_TEST_REDIS_ADDR")
|
||||
if addr == "" {
|
||||
t.Skip("set TELESRV_TEST_REDIS_ADDR to run redis integration test")
|
||||
|
|
@ -77,10 +61,10 @@ func TestRedisAllocatorConcurrentFirstUse(t *testing.T) {
|
|||
t.Cleanup(func() { _ = c.Close() })
|
||||
|
||||
userID := time.Now().UnixNano()
|
||||
t.Cleanup(func() { _ = c.Del(ctx, ptsKey(userID)).Err() })
|
||||
t.Cleanup(func() { _ = c.Del(ctx, boxIDKey(userID)).Err() })
|
||||
|
||||
const workers = 32
|
||||
pts := NewPtsAllocator(c, staticCounterSource{value: 1000})
|
||||
boxes := NewBoxIDAllocator(c, staticCounterSource{value: 1000})
|
||||
values := make(chan int, workers)
|
||||
errs := make(chan error, workers)
|
||||
var wg sync.WaitGroup
|
||||
|
|
@ -88,7 +72,7 @@ func TestRedisAllocatorConcurrentFirstUse(t *testing.T) {
|
|||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
v, err := pts.NextPts(ctx, userID)
|
||||
v, err := boxes.NextBoxID(ctx, userID)
|
||||
if err != nil {
|
||||
errs <- err
|
||||
return
|
||||
|
|
@ -101,29 +85,29 @@ func TestRedisAllocatorConcurrentFirstUse(t *testing.T) {
|
|||
close(errs)
|
||||
|
||||
for err := range errs {
|
||||
t.Fatalf("NextPts: %v", err)
|
||||
t.Fatalf("NextBoxID: %v", err)
|
||||
}
|
||||
seen := make(map[int]bool, workers)
|
||||
for v := range values {
|
||||
if v < 1001 || v > 1000+workers {
|
||||
t.Fatalf("pts = %d, want recovered contiguous range", v)
|
||||
t.Fatalf("box id = %d, want recovered contiguous range", v)
|
||||
}
|
||||
if seen[v] {
|
||||
t.Fatalf("duplicate pts %d", v)
|
||||
t.Fatalf("duplicate box id %d", v)
|
||||
}
|
||||
seen[v] = true
|
||||
}
|
||||
for want := 1001; want <= 1000+workers; want++ {
|
||||
if !seen[want] {
|
||||
t.Fatalf("missing pts %d", want)
|
||||
t.Fatalf("missing box id %d", want)
|
||||
}
|
||||
}
|
||||
current, err := pts.CurrentPts(ctx, userID)
|
||||
current, err := boxes.CurrentBoxID(ctx, userID)
|
||||
if err != nil {
|
||||
t.Fatalf("CurrentPts: %v", err)
|
||||
t.Fatalf("CurrentBoxID: %v", err)
|
||||
}
|
||||
if current != 1000+workers {
|
||||
t.Fatalf("current pts = %d, want %d", current, 1000+workers)
|
||||
t.Fatalf("current box id = %d, want %d", current, 1000+workers)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -157,4 +141,21 @@ func TestRedisRateLimiterWindow(t *testing.T) {
|
|||
if allowed || retry <= 0 {
|
||||
t.Fatalf("second allowed=%v retry=%d, want limited with retry", allowed, retry)
|
||||
}
|
||||
|
||||
batchKey := key + ":batch"
|
||||
t.Cleanup(func() { _ = c.Del(ctx, rateLimitKey(batchKey)).Err() })
|
||||
allowed, retry, err = limiter.AllowN(ctx, batchKey, 2, 3, time.Minute)
|
||||
if err != nil {
|
||||
t.Fatalf("AllowN first: %v", err)
|
||||
}
|
||||
if !allowed || retry != 0 {
|
||||
t.Fatalf("AllowN first allowed=%v retry=%d, want allowed", allowed, retry)
|
||||
}
|
||||
allowed, retry, err = limiter.AllowN(ctx, batchKey, 2, 3, time.Minute)
|
||||
if err != nil {
|
||||
t.Fatalf("AllowN second: %v", err)
|
||||
}
|
||||
if allowed || retry <= 0 {
|
||||
t.Fatalf("AllowN second allowed=%v retry=%d, want limited with retry", allowed, retry)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
363
internal/store/redisstore/inline_registry.go
Normal file
363
internal/store/redisstore/inline_registry.go
Normal file
|
|
@ -0,0 +1,363 @@
|
|||
package redisstore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
|
||||
"telesrv/internal/domain"
|
||||
"telesrv/internal/store"
|
||||
)
|
||||
|
||||
type InlineRegistryStore struct {
|
||||
c *redis.Client
|
||||
}
|
||||
|
||||
func NewInlineRegistryStore(c *redis.Client) *InlineRegistryStore {
|
||||
return &InlineRegistryStore{c: c}
|
||||
}
|
||||
|
||||
const inlineBotQueryChannel = "inline:bot_query"
|
||||
|
||||
func inlinePendingKey(queryID int64) string {
|
||||
return fmt.Sprintf("inline:pending:%d", queryID)
|
||||
}
|
||||
|
||||
func inlineResultKey(queryID int64) string {
|
||||
return fmt.Sprintf("inline:result:%d", queryID)
|
||||
}
|
||||
|
||||
func inlineCacheKey(key store.InlineCacheKey) (string, error) {
|
||||
raw, err := json.Marshal(key)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("marshal inline cache key: %w", err)
|
||||
}
|
||||
sum := sha256.Sum256(raw)
|
||||
return "inline:cache:" + hex.EncodeToString(sum[:]), nil
|
||||
}
|
||||
|
||||
func inlineWebDocumentKey(key store.InlineWebDocumentKey) string {
|
||||
sum := sha256.Sum256([]byte(fmt.Sprintf("%s\x00%d", key.URL, key.AccessHash)))
|
||||
return "inline:webdoc:" + hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func preparedInlineMessageKey(id string) string {
|
||||
sum := sha256.Sum256([]byte(id))
|
||||
return "inline:prepared:" + hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func webViewSessionKey(queryID int64) string {
|
||||
return fmt.Sprintf("webview:session:%d", queryID)
|
||||
}
|
||||
|
||||
func webViewBotQueryKey(botQueryID string) string {
|
||||
sum := sha256.Sum256([]byte(botQueryID))
|
||||
return "webview:bot_query:" + hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) PutInlinePending(ctx context.Context, pending store.InlinePending, ttl time.Duration) error {
|
||||
raw, err := json.Marshal(pending)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal inline pending: %w", err)
|
||||
}
|
||||
if err := s.c.Set(ctx, inlinePendingKey(pending.QueryID), raw, ttl).Err(); err != nil {
|
||||
return fmt.Errorf("redis set inline pending: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) GetInlinePending(ctx context.Context, queryID int64) (store.InlinePending, bool, error) {
|
||||
key := inlinePendingKey(queryID)
|
||||
var pending store.InlinePending
|
||||
found, err := redisGetJSON(ctx, s.c, key, &pending)
|
||||
if err != nil || !found {
|
||||
return store.InlinePending{}, false, err
|
||||
}
|
||||
if pending.QueryID != queryID || pending.BotUserID == 0 || pending.UserID == 0 {
|
||||
_ = s.c.Del(ctx, key).Err()
|
||||
return store.InlinePending{}, false, nil
|
||||
}
|
||||
return pending, true, nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) DeleteInlinePending(ctx context.Context, queryID int64) error {
|
||||
if err := s.c.Del(ctx, inlinePendingKey(queryID)).Err(); err != nil {
|
||||
return fmt.Errorf("redis delete inline pending: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) PutInlineResult(ctx context.Context, results domain.BotInlineResults, ttl time.Duration) error {
|
||||
raw, err := json.Marshal(results)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal inline result: %w", err)
|
||||
}
|
||||
if err := s.c.Set(ctx, inlineResultKey(results.QueryID), raw, ttl).Err(); err != nil {
|
||||
return fmt.Errorf("redis set inline result: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) GetInlineResult(ctx context.Context, queryID int64) (domain.BotInlineResults, bool, error) {
|
||||
key := inlineResultKey(queryID)
|
||||
var results domain.BotInlineResults
|
||||
found, err := redisGetJSON(ctx, s.c, key, &results)
|
||||
if err != nil || !found {
|
||||
return domain.BotInlineResults{}, false, err
|
||||
}
|
||||
if results.QueryID != queryID || results.UserID == 0 || results.BotUserID == 0 {
|
||||
_ = s.c.Del(ctx, key).Err()
|
||||
return domain.BotInlineResults{}, false, nil
|
||||
}
|
||||
return results, true, nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) DeleteInlineResult(ctx context.Context, queryID int64) error {
|
||||
if err := s.c.Del(ctx, inlineResultKey(queryID)).Err(); err != nil {
|
||||
return fmt.Errorf("redis delete inline result: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) PutInlineCache(ctx context.Context, key store.InlineCacheKey, results domain.BotInlineResults, ttl time.Duration) error {
|
||||
redisKey, err := inlineCacheKey(key)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
results.QueryID = 0
|
||||
raw, err := json.Marshal(results)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal inline cache: %w", err)
|
||||
}
|
||||
if err := s.c.Set(ctx, redisKey, raw, ttl).Err(); err != nil {
|
||||
return fmt.Errorf("redis set inline cache: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) GetInlineCache(ctx context.Context, key store.InlineCacheKey) (domain.BotInlineResults, bool, time.Duration, error) {
|
||||
redisKey, err := inlineCacheKey(key)
|
||||
if err != nil {
|
||||
return domain.BotInlineResults{}, false, 0, err
|
||||
}
|
||||
var results domain.BotInlineResults
|
||||
found, err := redisGetJSON(ctx, s.c, redisKey, &results)
|
||||
if err != nil || !found {
|
||||
return domain.BotInlineResults{}, false, 0, err
|
||||
}
|
||||
ttl, err := s.c.TTL(ctx, redisKey).Result()
|
||||
if err != nil {
|
||||
return domain.BotInlineResults{}, false, 0, fmt.Errorf("redis ttl inline cache: %w", err)
|
||||
}
|
||||
if ttl <= 0 {
|
||||
ttl = time.Second
|
||||
}
|
||||
return results, true, ttl, nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) PutInlineWebDocument(ctx context.Context, document domain.BotInlineWebDocument, ttl time.Duration) error {
|
||||
key := store.InlineWebDocumentKey{URL: document.URL, AccessHash: document.AccessHash}
|
||||
redisKey := inlineWebDocumentKey(key)
|
||||
entry := store.InlineWebDocumentEntry{Document: document}
|
||||
if existing, found, err := s.GetInlineWebDocument(ctx, key); err != nil {
|
||||
return err
|
||||
} else if found {
|
||||
entry.Bytes = append([]byte(nil), existing.Bytes...)
|
||||
entry.MimeType = existing.MimeType
|
||||
}
|
||||
raw, err := json.Marshal(entry)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal inline web document: %w", err)
|
||||
}
|
||||
if err := s.c.Set(ctx, redisKey, raw, ttl).Err(); err != nil {
|
||||
return fmt.Errorf("redis set inline web document: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) GetInlineWebDocument(ctx context.Context, key store.InlineWebDocumentKey) (store.InlineWebDocumentEntry, bool, error) {
|
||||
redisKey := inlineWebDocumentKey(key)
|
||||
var entry store.InlineWebDocumentEntry
|
||||
found, err := redisGetJSON(ctx, s.c, redisKey, &entry)
|
||||
if err != nil || !found {
|
||||
return store.InlineWebDocumentEntry{}, false, err
|
||||
}
|
||||
if entry.Document.URL != key.URL || entry.Document.AccessHash != key.AccessHash || entry.Document.URL == "" || entry.Document.AccessHash == 0 {
|
||||
_ = s.c.Del(ctx, redisKey).Err()
|
||||
return store.InlineWebDocumentEntry{}, false, nil
|
||||
}
|
||||
if len(entry.Bytes) > domain.MaxBotInlineWebSize {
|
||||
_ = s.c.Del(ctx, redisKey).Err()
|
||||
return store.InlineWebDocumentEntry{}, false, nil
|
||||
}
|
||||
entry.Bytes = append([]byte(nil), entry.Bytes...)
|
||||
return entry, true, nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) PutInlineWebDocumentBytes(ctx context.Context, key store.InlineWebDocumentKey, data []byte, mimeType string, ttl time.Duration) error {
|
||||
if len(data) == 0 || len(data) > domain.MaxBotInlineWebSize {
|
||||
return fmt.Errorf("inline web document bytes size %d out of range", len(data))
|
||||
}
|
||||
redisKey := inlineWebDocumentKey(key)
|
||||
entry, found, err := s.GetInlineWebDocument(ctx, key)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !found {
|
||||
return fmt.Errorf("inline web document missing")
|
||||
}
|
||||
entry.Bytes = append([]byte(nil), data...)
|
||||
entry.MimeType = mimeType
|
||||
raw, err := json.Marshal(entry)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal inline web document bytes: %w", err)
|
||||
}
|
||||
if currentTTL, err := s.c.TTL(ctx, redisKey).Result(); err == nil && currentTTL > 0 {
|
||||
ttl = currentTTL
|
||||
}
|
||||
if err := s.c.Set(ctx, redisKey, raw, ttl).Err(); err != nil {
|
||||
return fmt.Errorf("redis set inline web document bytes: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) PutPreparedInlineMessage(ctx context.Context, msg store.PreparedInlineMessage, ttl time.Duration) error {
|
||||
raw, err := json.Marshal(msg)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal prepared inline message: %w", err)
|
||||
}
|
||||
if err := s.c.Set(ctx, preparedInlineMessageKey(msg.ID), raw, ttl).Err(); err != nil {
|
||||
return fmt.Errorf("redis set prepared inline message: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) GetPreparedInlineMessage(ctx context.Context, id string) (store.PreparedInlineMessage, bool, error) {
|
||||
key := preparedInlineMessageKey(id)
|
||||
var msg store.PreparedInlineMessage
|
||||
found, err := redisGetJSON(ctx, s.c, key, &msg)
|
||||
if err != nil || !found {
|
||||
return store.PreparedInlineMessage{}, false, err
|
||||
}
|
||||
if msg.ID != id || msg.BotUserID == 0 || msg.UserID == 0 || len(msg.Results.Results) != 1 {
|
||||
_ = s.c.Del(ctx, key).Err()
|
||||
return store.PreparedInlineMessage{}, false, nil
|
||||
}
|
||||
msg.Results.Results = append([]domain.BotInlineResult(nil), msg.Results.Results...)
|
||||
return msg, true, nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) PutWebViewSession(ctx context.Context, session store.WebViewSession, ttl time.Duration) error {
|
||||
raw, err := json.Marshal(session)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal webview session: %w", err)
|
||||
}
|
||||
if err := s.c.Set(ctx, webViewSessionKey(session.QueryID), raw, ttl).Err(); err != nil {
|
||||
return fmt.Errorf("redis set webview session: %w", err)
|
||||
}
|
||||
if err := s.c.Set(ctx, webViewBotQueryKey(session.BotQueryID), raw, ttl).Err(); err != nil {
|
||||
return fmt.Errorf("redis set webview bot query: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) GetWebViewSession(ctx context.Context, queryID int64) (store.WebViewSession, bool, error) {
|
||||
key := webViewSessionKey(queryID)
|
||||
var session store.WebViewSession
|
||||
found, err := redisGetJSON(ctx, s.c, key, &session)
|
||||
if err != nil || !found {
|
||||
return store.WebViewSession{}, false, err
|
||||
}
|
||||
if !validWebViewSession(session) || session.QueryID != queryID {
|
||||
_ = s.c.Del(ctx, key).Err()
|
||||
return store.WebViewSession{}, false, nil
|
||||
}
|
||||
return session, true, nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) GetWebViewSessionByBotQuery(ctx context.Context, botQueryID string) (store.WebViewSession, bool, error) {
|
||||
key := webViewBotQueryKey(botQueryID)
|
||||
var session store.WebViewSession
|
||||
found, err := redisGetJSON(ctx, s.c, key, &session)
|
||||
if err != nil || !found {
|
||||
return store.WebViewSession{}, false, err
|
||||
}
|
||||
if !validWebViewSession(session) || session.BotQueryID != botQueryID {
|
||||
_ = s.c.Del(ctx, key).Err()
|
||||
return store.WebViewSession{}, false, nil
|
||||
}
|
||||
return session, true, nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) DeleteWebViewSession(ctx context.Context, queryID int64, botQueryID string) error {
|
||||
if err := s.c.Del(ctx, webViewSessionKey(queryID), webViewBotQueryKey(botQueryID)).Err(); err != nil {
|
||||
return fmt.Errorf("redis delete webview session: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validWebViewSession(session store.WebViewSession) bool {
|
||||
return session.QueryID != 0 && session.BotQueryID != "" && session.BotUserID != 0 && session.UserID != 0 && session.Peer.ID != 0
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) PublishBotInlineQuery(ctx context.Context, event store.BotInlineQueryPush) error {
|
||||
if event.SourceID == "" || event.QueryID == 0 || event.BotUserID == 0 || event.UserID == 0 {
|
||||
return fmt.Errorf("inline bot query push missing identity")
|
||||
}
|
||||
raw, err := json.Marshal(event)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal inline bot query push: %w", err)
|
||||
}
|
||||
if err := s.c.Publish(ctx, inlineBotQueryChannel, raw).Err(); err != nil {
|
||||
return fmt.Errorf("redis publish inline bot query: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *InlineRegistryStore) SubscribeBotInlineQueries(ctx context.Context, handle func(context.Context, store.BotInlineQueryPush)) error {
|
||||
if handle == nil {
|
||||
return fmt.Errorf("inline bot query handler is nil")
|
||||
}
|
||||
pubsub := s.c.Subscribe(ctx, inlineBotQueryChannel)
|
||||
defer func() { _ = pubsub.Close() }()
|
||||
if _, err := pubsub.Receive(ctx); err != nil {
|
||||
return fmt.Errorf("redis subscribe inline bot query: %w", err)
|
||||
}
|
||||
ch := pubsub.Channel()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case msg, ok := <-ch:
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
var event store.BotInlineQueryPush
|
||||
if err := json.Unmarshal([]byte(msg.Payload), &event); err != nil {
|
||||
continue
|
||||
}
|
||||
handle(ctx, event)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func redisGetJSON(ctx context.Context, c *redis.Client, key string, out any) (bool, error) {
|
||||
raw, err := c.Get(ctx, key).Bytes()
|
||||
if err != nil {
|
||||
if errors.Is(err, redis.Nil) {
|
||||
return false, nil
|
||||
}
|
||||
return false, fmt.Errorf("redis get %s: %w", key, err)
|
||||
}
|
||||
if err := json.Unmarshal(raw, out); err != nil {
|
||||
_ = c.Del(ctx, key).Err()
|
||||
return false, nil
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
|
@ -24,6 +24,13 @@ func rateLimitKey(key string) string {
|
|||
}
|
||||
|
||||
func (l *RateLimiter) Allow(ctx context.Context, key string, limit int, window time.Duration) (bool, int, error) {
|
||||
return l.AllowN(ctx, key, 1, limit, window)
|
||||
}
|
||||
|
||||
func (l *RateLimiter) AllowN(ctx context.Context, key string, cost, limit int, window time.Duration) (bool, int, error) {
|
||||
if cost <= 0 {
|
||||
return true, 0, nil
|
||||
}
|
||||
if limit <= 0 {
|
||||
return true, 0, nil
|
||||
}
|
||||
|
|
@ -34,11 +41,11 @@ func (l *RateLimiter) Allow(ctx context.Context, key string, limit int, window t
|
|||
return false, 0, fmt.Errorf("redis rate limiter: nil client")
|
||||
}
|
||||
redisKey := rateLimitKey(key)
|
||||
count, err := l.c.Incr(ctx, redisKey).Result()
|
||||
count, err := l.c.IncrBy(ctx, redisKey, int64(cost)).Result()
|
||||
if err != nil {
|
||||
return false, 0, fmt.Errorf("redis incr rate limit: %w", err)
|
||||
return false, 0, fmt.Errorf("redis incrby rate limit: %w", err)
|
||||
}
|
||||
if count == 1 {
|
||||
if count == int64(cost) {
|
||||
if err := l.c.Expire(ctx, redisKey, window).Err(); err != nil {
|
||||
return false, 0, fmt.Errorf("redis expire rate limit: %w", err)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -14,8 +14,9 @@ import (
|
|||
|
||||
// DefaultSessionTTL 是 session 记录的默认过期时间。
|
||||
// session 是连接态：过期或丢失后，客户端重连会触发 new_session_created / bad_server_salt 重建，
|
||||
// 因此 TTL 不必很长。
|
||||
const DefaultSessionTTL = 30 * 24 * time.Hour
|
||||
// 因此 TTL 不必很长。每个随机 session_id 都落一条记录且断连不删，过长的 TTL
|
||||
// 只会堆积死 session（移动端每次重连一条）。7 天足够覆盖常规离线窗口。
|
||||
const DefaultSessionTTL = 7 * 24 * time.Hour
|
||||
|
||||
// SessionStore 用 Redis 实现 store.SessionStore。
|
||||
type SessionStore struct {
|
||||
|
|
|
|||
235
internal/store/redisstore/user_cache.go
Normal file
235
internal/store/redisstore/user_cache.go
Normal file
|
|
@ -0,0 +1,235 @@
|
|||
package redisstore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
|
||||
"telesrv/internal/domain"
|
||||
)
|
||||
|
||||
const DefaultUserCacheTTL = 5 * time.Minute
|
||||
|
||||
// UserCache stores viewer-independent base user rows in Redis.
|
||||
type UserCache struct {
|
||||
c *redis.Client
|
||||
ttl time.Duration
|
||||
}
|
||||
|
||||
func NewUserCache(c *redis.Client, ttl time.Duration) *UserCache {
|
||||
if ttl <= 0 {
|
||||
ttl = DefaultUserCacheTTL
|
||||
}
|
||||
return &UserCache{c: c, ttl: ttl}
|
||||
}
|
||||
|
||||
func userBaseKey(id int64) string {
|
||||
return fmt.Sprintf("user:base:%d", id)
|
||||
}
|
||||
|
||||
type userBaseValue struct {
|
||||
ID int64 `json:"id"`
|
||||
AccessHash int64 `json:"access_hash"`
|
||||
Phone string `json:"phone"`
|
||||
FirstName string `json:"first_name"`
|
||||
LastName string `json:"last_name"`
|
||||
About string `json:"about"`
|
||||
Username string `json:"username"`
|
||||
CountryCode string `json:"country_code"`
|
||||
Verified bool `json:"verified"`
|
||||
Support bool `json:"support"`
|
||||
// bot 字段必须随缓存往返：丢失会让缓存命中路径把 bot 输出成普通用户，
|
||||
// 污染客户端本地缓存（TDesktop 的 bot 标记不可逆）。
|
||||
Bot bool `json:"bot,omitempty"`
|
||||
BotInfoVersion int `json:"bot_info_version,omitempty"`
|
||||
// premium / emoji status 同理必须随缓存往返：丢失会让缓存命中路径把
|
||||
// 会员输出成非会员，跨路径状态漂移（与 bot 列同一坑位）。
|
||||
PremiumUntil int `json:"premium_until,omitempty"`
|
||||
EmojiStatusDocumentID int64 `json:"emoji_status_document_id,omitempty"`
|
||||
EmojiStatusUntil int `json:"emoji_status_until,omitempty"`
|
||||
// birthday / personal channel 同理必须随缓存往返：缓存命中路径丢失会让刚保存的
|
||||
// 生日 / 个人频道在重新打开资料时归零（与 bot/premium 列同一坑位）。
|
||||
BirthdayDay int `json:"birthday_day,omitempty"`
|
||||
BirthdayMonth int `json:"birthday_month,omitempty"`
|
||||
BirthdayYear int `json:"birthday_year,omitempty"`
|
||||
PersonalChannelID int64 `json:"personal_channel_id,omitempty"`
|
||||
ColorSet bool `json:"color_set,omitempty"`
|
||||
Color int `json:"color,omitempty"`
|
||||
ColorBackgroundEmojiID int64 `json:"color_background_emoji_id,omitempty"`
|
||||
ProfileColorSet bool `json:"profile_color_set,omitempty"`
|
||||
ProfileColor int `json:"profile_color,omitempty"`
|
||||
ProfileColorBackgroundEmojiID int64 `json:"profile_color_background_emoji_id,omitempty"`
|
||||
LastSeenAt int `json:"last_seen_at"`
|
||||
}
|
||||
|
||||
func baseValueFromUser(u domain.User) userBaseValue {
|
||||
return userBaseValue{
|
||||
ID: u.ID,
|
||||
AccessHash: u.AccessHash,
|
||||
Phone: u.Phone,
|
||||
FirstName: u.FirstName,
|
||||
LastName: u.LastName,
|
||||
About: u.About,
|
||||
Username: u.Username,
|
||||
CountryCode: u.CountryCode,
|
||||
Verified: u.Verified,
|
||||
Support: u.Support,
|
||||
Bot: u.Bot,
|
||||
BotInfoVersion: u.BotInfoVersion,
|
||||
PremiumUntil: u.PremiumUntil,
|
||||
EmojiStatusDocumentID: u.EmojiStatusDocumentID,
|
||||
EmojiStatusUntil: u.EmojiStatusUntil,
|
||||
BirthdayDay: u.Birthday.Day,
|
||||
BirthdayMonth: u.Birthday.Month,
|
||||
BirthdayYear: u.Birthday.Year,
|
||||
PersonalChannelID: u.PersonalChannelID,
|
||||
ColorSet: u.Color.HasColor,
|
||||
Color: u.Color.Color,
|
||||
ColorBackgroundEmojiID: u.Color.BackgroundEmojiID,
|
||||
ProfileColorSet: u.ProfileColor.HasColor,
|
||||
ProfileColor: u.ProfileColor.Color,
|
||||
ProfileColorBackgroundEmojiID: u.ProfileColor.BackgroundEmojiID,
|
||||
LastSeenAt: u.LastSeenAt,
|
||||
}
|
||||
}
|
||||
|
||||
func (v userBaseValue) user() domain.User {
|
||||
return domain.User{
|
||||
ID: v.ID,
|
||||
AccessHash: v.AccessHash,
|
||||
Phone: v.Phone,
|
||||
FirstName: v.FirstName,
|
||||
LastName: v.LastName,
|
||||
About: v.About,
|
||||
Username: v.Username,
|
||||
CountryCode: v.CountryCode,
|
||||
Verified: v.Verified,
|
||||
Support: v.Support,
|
||||
Bot: v.Bot,
|
||||
BotInfoVersion: v.BotInfoVersion,
|
||||
PremiumUntil: v.PremiumUntil,
|
||||
EmojiStatusDocumentID: v.EmojiStatusDocumentID,
|
||||
EmojiStatusUntil: v.EmojiStatusUntil,
|
||||
Birthday: domain.Birthday{Day: v.BirthdayDay, Month: v.BirthdayMonth, Year: v.BirthdayYear},
|
||||
PersonalChannelID: v.PersonalChannelID,
|
||||
Color: domain.PeerColor{
|
||||
HasColor: v.ColorSet,
|
||||
Color: v.Color,
|
||||
BackgroundEmojiID: v.ColorBackgroundEmojiID,
|
||||
},
|
||||
ProfileColor: domain.PeerColor{
|
||||
HasColor: v.ProfileColorSet,
|
||||
Color: v.ProfileColor,
|
||||
BackgroundEmojiID: v.ProfileColorBackgroundEmojiID,
|
||||
},
|
||||
LastSeenAt: v.LastSeenAt,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *UserCache) GetByIDs(ctx context.Context, ids []int64) (map[int64]domain.User, error) {
|
||||
if s == nil || s.c == nil || len(ids) == 0 {
|
||||
return map[int64]domain.User{}, nil
|
||||
}
|
||||
unique := uniqueCacheUserIDs(ids)
|
||||
if len(unique) == 0 {
|
||||
return map[int64]domain.User{}, nil
|
||||
}
|
||||
keys := make([]string, 0, len(unique))
|
||||
for _, id := range unique {
|
||||
keys = append(keys, userBaseKey(id))
|
||||
}
|
||||
values, err := s.c.MGet(ctx, keys...).Result()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("redis mget user base: %w", err)
|
||||
}
|
||||
out := make(map[int64]domain.User, len(values))
|
||||
corruptKeys := make([]string, 0)
|
||||
for i, value := range values {
|
||||
if value == nil {
|
||||
continue
|
||||
}
|
||||
raw, ok := value.(string)
|
||||
if !ok {
|
||||
corruptKeys = append(corruptKeys, keys[i])
|
||||
continue
|
||||
}
|
||||
var decoded userBaseValue
|
||||
if err := json.Unmarshal([]byte(raw), &decoded); err != nil || decoded.ID == 0 || decoded.ID != unique[i] {
|
||||
corruptKeys = append(corruptKeys, keys[i])
|
||||
continue
|
||||
}
|
||||
out[decoded.ID] = decoded.user()
|
||||
}
|
||||
if len(corruptKeys) > 0 {
|
||||
_ = s.c.Del(ctx, corruptKeys...).Err()
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *UserCache) PutMany(ctx context.Context, users []domain.User) error {
|
||||
if s == nil || s.c == nil || len(users) == 0 {
|
||||
return nil
|
||||
}
|
||||
pipe := s.c.Pipeline()
|
||||
writes := 0
|
||||
seen := make(map[int64]struct{}, len(users))
|
||||
for _, u := range users {
|
||||
if u.ID == 0 {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[u.ID]; ok {
|
||||
continue
|
||||
}
|
||||
seen[u.ID] = struct{}{}
|
||||
raw, err := json.Marshal(baseValueFromUser(u))
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal user base cache: %w", err)
|
||||
}
|
||||
pipe.Set(ctx, userBaseKey(u.ID), raw, s.ttl)
|
||||
writes++
|
||||
}
|
||||
if writes == 0 {
|
||||
return nil
|
||||
}
|
||||
if _, err := pipe.Exec(ctx); err != nil {
|
||||
return fmt.Errorf("redis set user base: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *UserCache) Delete(ctx context.Context, ids []int64) error {
|
||||
if s == nil || s.c == nil || len(ids) == 0 {
|
||||
return nil
|
||||
}
|
||||
unique := uniqueCacheUserIDs(ids)
|
||||
if len(unique) == 0 {
|
||||
return nil
|
||||
}
|
||||
keys := make([]string, 0, len(unique))
|
||||
for _, id := range unique {
|
||||
keys = append(keys, userBaseKey(id))
|
||||
}
|
||||
if err := s.c.Del(ctx, keys...).Err(); err != nil {
|
||||
return fmt.Errorf("redis delete user base: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func uniqueCacheUserIDs(ids []int64) []int64 {
|
||||
out := make([]int64, 0, len(ids))
|
||||
seen := make(map[int64]struct{}, len(ids))
|
||||
for _, id := range ids {
|
||||
if id == 0 {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[id]; ok {
|
||||
continue
|
||||
}
|
||||
seen[id] = struct{}{}
|
||||
out = append(out, id)
|
||||
}
|
||||
return out
|
||||
}
|
||||
118
internal/store/redisstore/user_cache_integration_test.go
Normal file
118
internal/store/redisstore/user_cache_integration_test.go
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
package redisstore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"telesrv/internal/domain"
|
||||
)
|
||||
|
||||
func TestUserCacheRoundTrip(t *testing.T) {
|
||||
addr := os.Getenv("TELESRV_TEST_REDIS_ADDR")
|
||||
if addr == "" {
|
||||
t.Skip("set TELESRV_TEST_REDIS_ADDR to run redis integration test")
|
||||
}
|
||||
ctx := context.Background()
|
||||
c, err := Open(ctx, addr, "", 0)
|
||||
if err != nil {
|
||||
t.Fatalf("open: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = c.Close() })
|
||||
|
||||
const userID int64 = 77000001
|
||||
cache := NewUserCache(c, time.Minute)
|
||||
t.Cleanup(func() { _ = c.Del(ctx, userBaseKey(userID), userBaseKey(userID+1)).Err() })
|
||||
|
||||
want := domain.User{
|
||||
ID: userID,
|
||||
AccessHash: 12345,
|
||||
Phone: "15550000001",
|
||||
FirstName: "Alice",
|
||||
LastName: "Base",
|
||||
About: "about",
|
||||
Username: "alice_base",
|
||||
CountryCode: "US",
|
||||
Verified: true,
|
||||
Support: true,
|
||||
LastSeenAt: 99,
|
||||
Contact: true,
|
||||
PhotoID: 42,
|
||||
Birthday: domain.Birthday{Day: 14, Month: 2, Year: 1990},
|
||||
PersonalChannelID: 555,
|
||||
}
|
||||
if err := cache.PutMany(ctx, []domain.User{want, want}); err != nil {
|
||||
t.Fatalf("put: %v", err)
|
||||
}
|
||||
got, err := cache.GetByIDs(ctx, []int64{userID, userID + 1, userID})
|
||||
if err != nil {
|
||||
t.Fatalf("get: %v", err)
|
||||
}
|
||||
u, ok := got[userID]
|
||||
if !ok {
|
||||
t.Fatalf("cached user %d not found", userID)
|
||||
}
|
||||
if u.ID != want.ID || u.AccessHash != want.AccessHash || u.FirstName != want.FirstName || u.Username != want.Username || u.LastSeenAt != want.LastSeenAt {
|
||||
t.Fatalf("cached base mismatch: got %+v want %+v", u, want)
|
||||
}
|
||||
if u.Contact || u.PhotoID != 0 {
|
||||
t.Fatalf("viewer overlay leaked into base cache: %+v", u)
|
||||
}
|
||||
// birthday / personal channel 必须随缓存往返（缓存命中路径丢失会让刚保存的值归零）。
|
||||
if u.Birthday != want.Birthday || u.PersonalChannelID != want.PersonalChannelID {
|
||||
t.Fatalf("birthday/personal channel lost in base cache round-trip: got birthday=%+v personal=%d", u.Birthday, u.PersonalChannelID)
|
||||
}
|
||||
if _, ok := got[userID+1]; ok {
|
||||
t.Fatalf("unexpected missing user hit: %+v", got[userID+1])
|
||||
}
|
||||
|
||||
if err := c.Set(ctx, userBaseKey(userID+1), "{bad json", time.Minute).Err(); err != nil {
|
||||
t.Fatalf("set corrupt: %v", err)
|
||||
}
|
||||
got, err = cache.GetByIDs(ctx, []int64{userID + 1})
|
||||
if err != nil {
|
||||
t.Fatalf("get corrupt: %v", err)
|
||||
}
|
||||
if len(got) != 0 {
|
||||
t.Fatalf("corrupt cache returned users: %+v", got)
|
||||
}
|
||||
if n, err := c.Exists(ctx, userBaseKey(userID+1)).Result(); err != nil || n != 0 {
|
||||
t.Fatalf("corrupt key exists=%d err=%v, want deleted", n, err)
|
||||
}
|
||||
|
||||
if err := cache.PutMany(ctx, []domain.User{{ID: userID + 1, AccessHash: 6789, FirstName: "WrongKey"}}); err != nil {
|
||||
t.Fatalf("put mismatched payload source: %v", err)
|
||||
}
|
||||
raw, err := c.Get(ctx, userBaseKey(userID+1)).Result()
|
||||
if err != nil {
|
||||
t.Fatalf("get raw mismatched source: %v", err)
|
||||
}
|
||||
if err := c.Set(ctx, userBaseKey(userID), raw, time.Minute).Err(); err != nil {
|
||||
t.Fatalf("set mismatched payload: %v", err)
|
||||
}
|
||||
got, err = cache.GetByIDs(ctx, []int64{userID})
|
||||
if err != nil {
|
||||
t.Fatalf("get mismatched payload: %v", err)
|
||||
}
|
||||
if len(got) != 0 {
|
||||
t.Fatalf("mismatched payload returned users: %+v", got)
|
||||
}
|
||||
if n, err := c.Exists(ctx, userBaseKey(userID)).Result(); err != nil || n != 0 {
|
||||
t.Fatalf("mismatched key exists=%d err=%v, want deleted", n, err)
|
||||
}
|
||||
|
||||
if err := cache.PutMany(ctx, []domain.User{want}); err != nil {
|
||||
t.Fatalf("restore user: %v", err)
|
||||
}
|
||||
if err := cache.Delete(ctx, []int64{userID}); err != nil {
|
||||
t.Fatalf("delete: %v", err)
|
||||
}
|
||||
got, err = cache.GetByIDs(ctx, []int64{userID})
|
||||
if err != nil {
|
||||
t.Fatalf("get after delete: %v", err)
|
||||
}
|
||||
if len(got) != 0 {
|
||||
t.Fatalf("cache hit after delete: %+v", got)
|
||||
}
|
||||
}
|
||||
|
|
@ -1,112 +0,0 @@
|
|||
package redisstore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
|
||||
"telesrv/internal/store"
|
||||
)
|
||||
|
||||
// UserCounterAllocator 在一次 Lua 往返内同时分配某 user 的下一个 pts 与 box_id。
|
||||
//
|
||||
// counter:pts:{userID} 与 counter:box_id:{userID} 共享 hash-tag {userID}，处于同一
|
||||
// Redis Cluster slot，可被单个 Lua 脚本原子操作——把发送热路径上「pts + box 两次往返」
|
||||
// 合并成一次。语义与 PtsAllocator/BoxIDAllocator 各自调用等价：pts 账号级无洞、
|
||||
// Redis 冷 miss 时从 PG durable log 恢复基线。
|
||||
type UserCounterAllocator struct {
|
||||
c *redis.Client
|
||||
ptsSource store.CounterSource
|
||||
boxSource store.CounterSource
|
||||
}
|
||||
|
||||
// NewUserCounterAllocator 创建合并 allocator。ptsSource 恢复 pts（MAX(user_update_events.pts)），
|
||||
// boxSource 恢复 box_id（MAX(message_boxes.box_id)）。
|
||||
func NewUserCounterAllocator(c *redis.Client, ptsSource, boxSource store.CounterSource) *UserCounterAllocator {
|
||||
return &UserCounterAllocator{c: c, ptsSource: ptsSource, boxSource: boxSource}
|
||||
}
|
||||
|
||||
// userCountersNextScript 热路径：两 key 都存在时各自 INCR 并返回；任一缺失返回 {-1,-1} 触发恢复。
|
||||
var userCountersNextScript = redis.NewScript(`
|
||||
local pts = redis.call("GET", KEYS[1])
|
||||
local box = redis.call("GET", KEYS[2])
|
||||
if pts and box then
|
||||
return {redis.call("INCR", KEYS[1]), redis.call("INCR", KEYS[2])}
|
||||
end
|
||||
return {-1, -1}
|
||||
`)
|
||||
|
||||
// userCountersRecoverScript 冷路径：每个 key「缺失才用 PG 基线 SET，再 INCR」。
|
||||
// 对已存在的 key 跳过 SET 只 INCR，故对「一存一缺」也安全；Redis 单线程串行化保证无重复无洞。
|
||||
var userCountersRecoverScript = redis.NewScript(`
|
||||
local function recnext(key, base)
|
||||
if not redis.call("GET", key) then
|
||||
redis.call("SET", key, base)
|
||||
end
|
||||
return redis.call("INCR", key)
|
||||
end
|
||||
return {recnext(KEYS[1], ARGV[1]), recnext(KEYS[2], ARGV[2])}
|
||||
`)
|
||||
|
||||
// NextUserCounters 返回该 user 的下一个 (pts, boxID)。
|
||||
func (a *UserCounterAllocator) NextUserCounters(ctx context.Context, userID int64) (int, int, error) {
|
||||
if userID == 0 {
|
||||
return 0, 0, fmt.Errorf("redis user counters: missing user id")
|
||||
}
|
||||
if a.c == nil {
|
||||
return 0, 0, fmt.Errorf("redis user counters: nil client")
|
||||
}
|
||||
keys := []string{ptsKey(userID), boxIDKey(userID)}
|
||||
|
||||
pts, box, err := runTwoCounters(ctx, a.c, userCountersNextScript, keys)
|
||||
if err != nil {
|
||||
return 0, 0, fmt.Errorf("redis next user counters: %w", err)
|
||||
}
|
||||
if pts != missingCounterSentinel && box != missingCounterSentinel {
|
||||
return int(pts), int(box), nil
|
||||
}
|
||||
|
||||
// 冷路径：至少一个 key 缺失，从 PG durable log 恢复两个基线，再 recover-next。
|
||||
ptsBase, err := recoverBase(ctx, a.ptsSource, userID, "pts")
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
boxBase, err := recoverBase(ctx, a.boxSource, userID, "box_id")
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
pts, box, err = runTwoCounters(ctx, a.c, userCountersRecoverScript, keys, ptsBase, boxBase)
|
||||
if err != nil {
|
||||
return 0, 0, fmt.Errorf("redis recover user counters: %w", err)
|
||||
}
|
||||
return int(pts), int(box), nil
|
||||
}
|
||||
|
||||
// runTwoCounters 执行返回二元整数表的 Lua 脚本，用 .Slice() 手动解析（不依赖 Int64Slice）。
|
||||
func runTwoCounters(ctx context.Context, c *redis.Client, script *redis.Script, keys []string, args ...any) (int64, int64, error) {
|
||||
raw, err := script.Run(ctx, c, keys, args...).Slice()
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
if len(raw) != 2 {
|
||||
return 0, 0, fmt.Errorf("unexpected reply len %d, want 2", len(raw))
|
||||
}
|
||||
a, okA := raw[0].(int64)
|
||||
b, okB := raw[1].(int64)
|
||||
if !okA || !okB {
|
||||
return 0, 0, fmt.Errorf("unexpected reply element types %T,%T, want int64", raw[0], raw[1])
|
||||
}
|
||||
return a, b, nil
|
||||
}
|
||||
|
||||
func recoverBase(ctx context.Context, source store.CounterSource, userID int64, name string) (int, error) {
|
||||
if source == nil {
|
||||
return 0, nil
|
||||
}
|
||||
v, err := source.Current(ctx, userID)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("recover %s counter: %w", name, err)
|
||||
}
|
||||
return v, nil
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue