fix: rate limit ephemeral room picks before querying upstream

internal/adapter/in/api/catalog_test.go

@@ -46,6 +46,10 @@ func (q *quota) TakeSearch(context.Context, string, int, time.Duration) (time.Du
 	return 0, nil
 }
 
+func (q *quota) TakePick(context.Context, string, int, time.Duration) (time.Duration, error) {
+	return 0, nil
+}
+
 var rain = port.Song{
 	Track: track.Track{Source: track.NetEase, ID: "185811"},
 	Meta:  track.Meta{Title: "七里香", Artist: "周杰伦", Album: "七里香", PageURL: "https://music.163.com/#/song?id=185811", Duration: 299 * time.Second},

internal/adapter/in/api/convert.go

@@ -25,12 +25,12 @@ var roomErrorCodes = map[room.ErrorKind]protocol.ErrorCode{
 }
 
 // bot 表示 by 是常驻房间的 bot 成员。返回的成员是房间命令的发起人，bot 代填了点歌人时换成那个名字。
-func (s *server) action(ctx context.Context, by member.Member, bot bool, data any) (room.Action, member.Member, error) {
+func (s *server) action(ctx context.Context, roomID string, by member.Member, bot bool, data any) (room.Action, member.Member, error) {
 	switch d := data.(type) {
 	case protocol.RoomSyncRequest:
 		return room.Sync{}, by, nil
 	case protocol.QueueAddRequest:
-		return s.Gateway.Resolve(ctx, by, bot, pick(d))
+		return s.Gateway.Resolve(ctx, roomID, by, bot, pick(d))
 	case protocol.QueueRemoveRequest:
 		return room.Remove{Entry: entryID(d.EntryID)}, by, nil
 	case protocol.QueueLikeRequest:

internal/adapter/in/api/room.go

@@ -28,7 +28,7 @@ func (s *server) roomRequest(route protocol.RoomRoute) handlerFunc {
 		if resp, handled, err := s.answer(ctx, c.roomID, data); handled {
 			return resp, err
 		}
-		a, by, err := s.action(ctx, c.member, c.bot, data)
+		a, by, err := s.action(ctx, c.roomID, c.member, c.bot, data)
 		if err != nil {
 			return nil, err
 		}

internal/adapter/in/api/ws.go

@@ -198,7 +198,7 @@ func (s *server) perform(ctx context.Context, sess *realtime.Session, tag reques
 		return err
 	}
 	// 握手拒绝了 API Key，会话的成员不是 bot，房间命令总以会话的成员发出。
-	a, _, err := s.action(ctx, sess.Member(), false, data)
+	a, _, err := s.action(ctx, sess.Room(), sess.Member(), false, data)
 	if err != nil {
 		sess.Fail(tag, err)
 		return err

internal/adapter/out/redis/quota.go

@@ -12,20 +12,28 @@ import (
 	"github.com/dihemoe/alisten/internal/app/port"
 )
 
-var _ port.SearchQuota = (*SearchQuota)(nil)
+var _ port.RoomQuota = (*RoomQuota)(nil)
 
-type SearchQuota struct {
+type RoomQuota struct {
 	limiter *redis_rate.Limiter
 }
 
-func NewSearchQuota(client *goredis.Client) *SearchQuota {
-	return &SearchQuota{limiter: redis_rate.NewLimiter(client)}
+func NewRoomQuota(client *goredis.Client) *RoomQuota {
+	return &RoomQuota{limiter: redis_rate.NewLimiter(client)}
 }
 
-func (q *SearchQuota) TakeSearch(ctx context.Context, roomID string, burst int, refill time.Duration) (time.Duration, error) {
-	res, err := q.limiter.Allow(ctx, "room:"+roomID+":search", redis_rate.Limit{Rate: 1, Burst: burst, Period: refill})
+func (q *RoomQuota) TakeSearch(ctx context.Context, roomID string, burst int, refill time.Duration) (time.Duration, error) {
+	return q.take(ctx, roomID, "search", "搜索", burst, refill)
+}
+
+func (q *RoomQuota) TakePick(ctx context.Context, roomID string, burst int, refill time.Duration) (time.Duration, error) {
+	return q.take(ctx, roomID, "pick", "点歌", burst, refill)
+}
+
+func (q *RoomQuota) take(ctx context.Context, roomID, bucket, what string, burst int, refill time.Duration) (time.Duration, error) {
+	res, err := q.limiter.Allow(ctx, "room:"+roomID+":"+bucket, redis_rate.Limit{Rate: 1, Burst: burst, Period: refill})
 	if err != nil {
-		return 0, fmt.Errorf("取房间 %s 的搜索额度：%w", roomID, err)
+		return 0, fmt.Errorf("取房间 %s 的%s额度：%w", roomID, what, err)
 	}
 	if res.Allowed > 0 {
 		return 0, nil

internal/adapter/out/redis/quota_test.go

@@ -6,7 +6,7 @@ import (
 )
 
 func TestTakeSearch(t *testing.T) {
-	quota := NewSearchQuota(newClient(t))
+	quota := NewRoomQuota(newClient(t))
 	for i := range 2 {
 		if wait, err := quota.TakeSearch(t.Context(), "room-1", 2, time.Minute); err != nil || wait != 0 {
 			t.Fatalf("第 %d 次 = %v, %v，突发额度之内应直接取到", i+1, wait, err)
@@ -23,3 +23,23 @@ func TestTakeSearch(t *testing.T) {
 		t.Errorf("另一个房间 = %v, %v，额度按房间计算", wait, err)
 	}
 }
+
+func TestTakePick(t *testing.T) {
+	quota := NewRoomQuota(newClient(t))
+	if wait, err := quota.TakeSearch(t.Context(), "room-1", 1, time.Minute); err != nil || wait != 0 {
+		t.Fatalf("搜索 = %v, %v", wait, err)
+	}
+	if wait, err := quota.TakePick(t.Context(), "room-1", 1, time.Minute); err != nil || wait != 0 {
+		t.Fatalf("搜索额度用完之后点歌 = %v, %v，点歌额度与搜索额度分开计算", wait, err)
+	}
+	wait, err := quota.TakePick(t.Context(), "room-1", 1, time.Minute)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if wait < 59*time.Second || wait > time.Minute {
+		t.Errorf("点歌额度用完后要等 %v，期望约 1 分钟", wait)
+	}
+	if wait, err := quota.TakePick(t.Context(), "room-2", 1, time.Minute); err != nil || wait != 0 {
+		t.Errorf("另一个房间 = %v, %v，点歌额度按房间计算", wait, err)
+	}
+}

internal/app/catalog/catalog.go

@@ -39,10 +39,10 @@ type Service struct {
 	tracks   port.Tracks
 	history  port.PlayHistory
 	rooms    port.RoomDirectory
-	quota    port.SearchQuota
+	quota    port.RoomQuota
 }
 
-func New(netease port.NetEaseSource, qq port.Source, bilibili port.BilibiliSource, links port.BilibiliLinks, tracks port.Tracks, history port.PlayHistory, rooms port.RoomDirectory, quota port.SearchQuota) *Service {
+func New(netease port.NetEaseSource, qq port.Source, bilibili port.BilibiliSource, links port.BilibiliLinks, tracks port.Tracks, history port.PlayHistory, rooms port.RoomDirectory, quota port.RoomQuota) *Service {
 	return &Service{netease: netease, qq: qq, bilibili: bilibili, links: links, tracks: tracks, history: history, rooms: rooms, quota: quota}
 }
 
@@ -209,15 +209,27 @@ func (s *Service) candidates(ctx context.Context, roomID string, includeQueue bo
 	return picks[:min(len(picks), autoDJLimit)], nil
 }
 
-func (s *Service) take(ctx context.Context, roomID string) error {
-	readCtx, cancel := context.WithTimeout(ctx, storeTimeout)
-	r, err := s.rooms.FindRoom(readCtx, roomID)
-	cancel()
+// 点歌要先向上游查询曲目，房间的点歌限流却在收到命令之后才判断，所以查询之前按同样的规则预扣一次，
+// 房间仍按自己的限流做最终裁决。
+func (s *Service) TakePick(ctx context.Context, roomID string) error {
+	ephemeral, err := s.ephemeral(ctx, roomID)
+	if err != nil || !ephemeral {
+		return err
+	}
+	wait, err := s.quota.TakePick(ctx, roomID, room.EphemeralBurst, room.RefillInterval)
 	if err != nil {
 		return err
 	}
-	if r.Info.Kind == room.Persistent {
-		return nil
+	if wait > 0 {
+		return room.ErrRateLimited
+	}
+	return nil
+}
+
+func (s *Service) take(ctx context.Context, roomID string) error {
+	ephemeral, err := s.ephemeral(ctx, roomID)
+	if err != nil || !ephemeral {
+		return err
 	}
 	for {
 		wait, err := s.quota.TakeSearch(ctx, roomID, searchBurst, searchRefill)
@@ -240,6 +252,16 @@ func (s *Service) take(ctx context.Context, roomID string) error {
 	}
 }
 
+func (s *Service) ephemeral(ctx context.Context, roomID string) (bool, error) {
+	readCtx, cancel := context.WithTimeout(ctx, storeTimeout)
+	defer cancel()
+	r, err := s.rooms.FindRoom(readCtx, roomID)
+	if err != nil {
+		return false, err
+	}
+	return r.Info.Kind != room.Persistent, nil
+}
+
 func (s *Service) metaSource(source track.Source) (func(context.Context, string) (track.Meta, error), error) {
 	switch source {
 	case track.NetEase:

internal/app/catalog/catalog_test.go

@@ -44,9 +44,11 @@ func (h history) AutoDJExclusions(context.Context, string) ([]track.Track, error
 }
 
 type quota struct {
-	mu    sync.Mutex
-	waits []time.Duration
-	calls int
+	mu        sync.Mutex
+	waits     []time.Duration
+	calls     int
+	pickWait  time.Duration
+	pickCalls int
 }
 
 func (q *quota) TakeSearch(_ context.Context, _ string, burst int, refill time.Duration) (time.Duration, error) {
@@ -64,6 +66,16 @@ func (q *quota) TakeSearch(_ context.Context, _ string, burst int, refill time.D
 	return wait, nil
 }
 
+func (q *quota) TakePick(_ context.Context, _ string, burst int, refill time.Duration) (time.Duration, error) {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	if burst != 5 || refill != time.Minute {
+		return 0, errors.New("临时房间的点歌额度与房间的点歌限流相同：突发 5 次，之后每分钟恢复 1 次")
+	}
+	q.pickCalls++
+	return q.pickWait, nil
+}
+
 type fixture struct {
 	svc        *Service
 	netease    *porttest.Source
@@ -174,6 +186,32 @@ func TestSearchQuota(t *testing.T) {
 	})
 }
 
+func TestTakePick(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		f := newFixture(t, history{})
+		if err := f.svc.TakePick(t.Context(), f.ephemeral); err != nil || f.quota.pickCalls != 1 {
+			t.Fatalf("临时房间 = %v，取了 %d 次点歌额度，期望取到 1 次", err, f.quota.pickCalls)
+		}
+
+		f.quota.pickWait = time.Minute
+		start := time.Now()
+		if err := f.svc.TakePick(t.Context(), f.ephemeral); !errors.Is(err, room.ErrRateLimited) {
+			t.Errorf("点歌额度用完时的错误 = %v，期望 room.ErrRateLimited", err)
+		}
+		if waited := time.Since(start); waited != 0 {
+			t.Errorf("等了 %v，期望与房间的点歌限流一样立即拒绝", waited)
+		}
+
+		calls := f.quota.pickCalls
+		if err := f.svc.TakePick(t.Context(), f.persistent); err != nil || f.quota.pickCalls != calls {
+			t.Errorf("常驻房间 = %v，取了 %d 次点歌额度，期望不取", err, f.quota.pickCalls-calls)
+		}
+		if f.quota.calls != 0 {
+			t.Errorf("取了 %d 次搜索额度，点歌只取点歌额度", f.quota.calls)
+		}
+	})
+}
+
 func TestTrackMetaPrefersTracksTable(t *testing.T) {
 	f := newFixture(t, history{})
 	sunny := song(track.NetEase, "1", "晴天")

internal/app/port/quota.go

@@ -5,9 +5,11 @@ import (
 	"time"
 )
 
-// 搜索与推荐列表的额度按房间计数，所有实例共用。
-type SearchQuota interface {
+// 额度按房间计数，所有实例共用。搜索与推荐列表共用一份额度，点歌另算一份。
+type RoomQuota interface {
 	// 从房间的额度中取一次：最多突发 burst 次，之后每隔 refill 恢复一次。取到时返回 0，
 	// 额度用完时返回距离下一次有额度还要等待的时长。
 	TakeSearch(ctx context.Context, roomID string, burst int, refill time.Duration) (time.Duration, error)
+	// 与 TakeSearch 相同，取的是点歌额度。
+	TakePick(ctx context.Context, roomID string, burst int, refill time.Duration) (time.Duration, error)
 }

internal/app/realtime/gateway.go

@@ -277,13 +277,16 @@ type Pick struct {
 }
 
 // by 是发出请求的成员，bot 表示他是常驻房间的 bot 成员。返回的成员是条目的点歌人，房间命令以他的名义发出。
-func (g *Gateway) Resolve(ctx context.Context, by member.Member, bot bool, p Pick) (room.Enqueue, member.Member, error) {
+func (g *Gateway) Resolve(ctx context.Context, roomID string, by member.Member, bot bool, p Pick) (room.Enqueue, member.Member, error) {
 	if p.Requester != "" {
 		if !bot {
 			return room.Enqueue{}, member.Member{}, ErrRequesterNotAllowed
 		}
 		by = member.Member{ID: by.ID, Name: p.Requester + "（bot）"}
 	}
+	if err := g.catalog.TakePick(ctx, roomID); err != nil {
+		return room.Enqueue{}, member.Member{}, err
+	}
 	switch {
 	case p.Keyword != "":
 		t, m, err := g.catalog.FirstMatch(ctx, p.Source, p.Keyword)

internal/app/realtime/gateway_test.go

@@ -8,6 +8,7 @@ import (
 	"slices"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"testing/synctest"
 	"time"
@@ -88,14 +89,32 @@ type gatewayFixture struct {
 	store *porttest.Store
 	redis *porttest.Redis
 	songs *songs
+	quota *quota
 }
 
 func newGatewayFixture(t *testing.T, snapshots port.RoomSnapshots, store *porttest.Store) *gatewayFixture {
 	t.Helper()
-	songs := newSongs()
+	q := &quota{}
+	songs := newPickingSongs(store, q)
 	songs.add(sunny, sunnyM, sunnyP)
 	redis := porttest.NewRedis()
-	return &gatewayFixture{instance: startInstance(t, redis, "a", snapshots, store, songs), store: store, redis: redis, songs: songs}
+	return &gatewayFixture{instance: startInstance(t, redis, "a", snapshots, store, songs), store: store, redis: redis, songs: songs, quota: q}
+}
+
+// quota 的点歌额度可以耗尽，耗尽之后要等一分钟；搜索额度总能取到。
+type quota struct {
+	exhausted atomic.Bool
+}
+
+func (q *quota) TakeSearch(context.Context, string, int, time.Duration) (time.Duration, error) {
+	return 0, nil
+}
+
+func (q *quota) TakePick(context.Context, string, int, time.Duration) (time.Duration, error) {
+	if q.exhausted.Load() {
+		return time.Minute, nil
+	}
+	return 0, nil
 }
 
 func createRoom(t *testing.T, store *porttest.Store, kind room.Kind) string {
@@ -375,6 +394,7 @@ func TestResolvePick(t *testing.T) {
 	synctest.Test(t, func(t *testing.T) {
 		store := porttest.New()
 		f := newGatewayFixture(t, store, store)
+		id := createRoom(t, store, room.Ephemeral)
 		link := track.Track{Source: track.Link, ID: "https://www.youtube.com/watch?v=x"}
 		video := track.Track{Source: track.Bilibili, ID: "BV1xx411c7mD"}
 		videoM := track.Meta{Title: "晴天（钢琴版）", PageURL: "https://www.bilibili.com/video/BV1xx411c7mD"}
@@ -392,7 +412,7 @@ func TestResolvePick(t *testing.T) {
 			{Pick{Source: track.NetEase, Keyword: "晴"}, room.Enqueue{Track: sunny, Meta: sunnyM}, nil},
 		}
 		for _, tt := range tests {
-			got, _, err := f.gateway.Resolve(t.Context(), alice, false, tt.pick)
+			got, _, err := f.gateway.Resolve(t.Context(), id, alice, false, tt.pick)
 			if got != tt.want || !errors.Is(err, tt.err) {
 				t.Errorf("Resolve(%+v) = %+v, %v，期望 %+v, %v", tt.pick, got, err, tt.want, tt.err)
 			}
@@ -404,19 +424,38 @@ func TestOnlyBotsFillInRequester(t *testing.T) {
 	synctest.Test(t, func(t *testing.T) {
 		store := porttest.New()
 		f := newGatewayFixture(t, store, store)
+		id := createRoom(t, store, room.Persistent)
 		bot := member.Member{ID: "00000000-0000-7000-8000-0000000000b0", Name: "bot"}
 		for requester, want := range map[string]member.Member{"张三": {ID: bot.ID, Name: "张三（bot）"}, "": bot} {
-			got, by, err := f.gateway.Resolve(t.Context(), bot, true, Pick{Track: sunny, Requester: requester})
+			got, by, err := f.gateway.Resolve(t.Context(), id, bot, true, Pick{Track: sunny, Requester: requester})
 			if err != nil || got.Track != sunny || by != want {
 				t.Errorf("bot 点歌时代填 %q = %+v, %+v, %v，期望点歌人 %+v", requester, got, by, err, want)
 			}
 		}
-		if _, _, err := f.gateway.Resolve(t.Context(), alice, false, Pick{Track: sunny, Requester: "张三"}); !errors.Is(err, ErrRequesterNotAllowed) {
+		if _, _, err := f.gateway.Resolve(t.Context(), id, alice, false, Pick{Track: sunny, Requester: "张三"}); !errors.Is(err, ErrRequesterNotAllowed) {
 			t.Errorf("普通成员代填点歌人的错误 = %v，期望 ErrRequesterNotAllowed", err)
 		}
 	})
 }
 
+func TestPickQuotaBeforeUpstream(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		store := porttest.New()
+		f := newGatewayFixture(t, store, store)
+		ephemeral, persistent := createRoom(t, store, room.Ephemeral), createRoom(t, store, room.Persistent)
+		f.quota.exhausted.Store(true)
+		f.songs.netease.Fail(errors.New("上游不可用"))
+		for _, p := range []Pick{{Track: sunny}, {Source: track.NetEase, Keyword: "晴"}} {
+			if _, _, err := f.gateway.Resolve(t.Context(), ephemeral, alice, false, p); !errors.Is(err, room.ErrRateLimited) {
+				t.Errorf("临时房间的点歌额度用完时 Resolve(%+v) 的错误 = %v，期望查询上游之前就以 room.ErrRateLimited 拒绝", p, err)
+			}
+		}
+		if _, _, err := f.gateway.Resolve(t.Context(), persistent, alice, false, Pick{Track: sunny}); !errors.Is(err, port.ErrProviderFailed) {
+			t.Errorf("常驻房间 Resolve 的错误 = %v，期望不预扣点歌额度，直接查询上游", err)
+		}
+	})
+}
+
 func TestUnsubscribeAfterLastUse(t *testing.T) {
 	synctest.Test(t, func(t *testing.T) {
 		store := porttest.New()

internal/app/realtime/host_test.go

@@ -66,10 +66,14 @@ type songs struct {
 	preparer *porttest.Preparer
 }
 
+// 只有网关点歌（Resolve）要按房间预扣点歌额度，用到房间目录和额度；不点歌的测试用 newSongs。
 func newSongs() *songs {
+	return newPickingSongs(nil, nil)
+}
+
+func newPickingSongs(rooms port.RoomDirectory, quota port.RoomQuota) *songs {
 	netease, bilibili := porttest.NewSource(track.NetEase), porttest.NewBilibili()
-	// 网关和房间宿主只用到 TrackMeta、ResolveLink、FirstMatch 和 Recommend，用不到房间目录和搜索额度。
-	music := catalog.New(netease, porttest.NewSource(track.QQ), bilibili, bilibili, noTracks{}, noHistory{}, nil, nil)
+	music := catalog.New(netease, porttest.NewSource(track.QQ), bilibili, bilibili, noTracks{}, noHistory{}, rooms, quota)
 	return &songs{netease: netease, bilibili: bilibili, music: music, preparer: porttest.NewPreparer()}
 }
 

internal/bootstrap/public.go

@@ -64,7 +64,7 @@ func newPublic(cfg serveConfig, pool *pgxpool.Pool, rdb *goredis.Client, instanc
 		return public{}, fmt.Errorf("创建 River client：%w", err)
 	}
 	preparer := prepare.New(media.NewTracks(store, objects), queue.New(jobs), logger)
-	music := catalog.New(src.netease, src.qq, src.bilibili, src.bilibiliLinks, store, store, store, redis.NewSearchQuota(rdb))
+	music := catalog.New(src.netease, src.qq, src.bilibili, src.bilibiliLinks, store, store, store, redis.NewRoomQuota(rdb))
 	leases := redis.NewLeases(rdb, instance)
 	cluster := realtime.Cluster{
 		Instance: instance,

internal/domain/room/error.go

@@ -21,7 +21,7 @@ type Error struct {
 func (e *Error) Error() string { return e.Message }
 
 var (
-	errRateLimited    = &Error{RateLimited, "操作过于频繁，请稍后再试。"}
+	ErrRateLimited    = &Error{RateLimited, "操作过于频繁，请稍后再试。"}
 	errLinkForbidden  = &Error{Forbidden, "临时房间只能点哔哩哔哩的链接。"}
 	errQueueFull      = &Error{QueueFull, fmt.Sprintf("临时房间的播放列表最多 %d 首歌，等播放几首之后再点。", ephemeralQueueLimit)}
 	errDuplicatePick  = &Error{Conflict, "这首歌已经在播放列表里了。"}

internal/domain/room/queue.go

@@ -37,12 +37,12 @@ func (r *Room) burst() int {
 	if r.info.Kind == Persistent {
 		return persistentBurst
 	}
-	return ephemeralBurst
+	return EphemeralBurst
 }
 
 func (r *Room) enqueue(b *batch, by member.Member, a Enqueue) (Result, *Error) {
-	if !r.pickLimit.allow(b.now, r.burst(), refillInterval) {
-		return nil, errRateLimited
+	if !r.pickLimit.allow(b.now, r.burst(), RefillInterval) {
+		return nil, ErrRateLimited
 	}
 	if a.Track.Source == track.Link && r.info.Kind == Ephemeral {
 		return nil, errLinkForbidden
@@ -61,8 +61,8 @@ func (r *Room) enqueue(b *batch, by member.Member, a Enqueue) (Result, *Error) {
 }
 
 func (r *Room) remove(b *batch, by member.Member, a Remove) (Result, *Error) {
-	if !r.pickLimit.allow(b.now, r.burst(), refillInterval) {
-		return nil, errRateLimited
+	if !r.pickLimit.allow(b.now, r.burst(), RefillInterval) {
+		return nil, ErrRateLimited
 	}
 	i := r.waitingIndex(a.Entry)
 	if i < 0 {
@@ -75,8 +75,8 @@ func (r *Room) remove(b *batch, by member.Member, a Remove) (Result, *Error) {
 }
 
 func (r *Room) like(b *batch, by member.Member, a Like) (Result, *Error) {
-	if !r.likeLimit.allow(b.now, r.burst(), refillInterval) {
-		return nil, errRateLimited
+	if !r.likeLimit.allow(b.now, r.burst(), RefillInterval) {
+		return nil, ErrRateLimited
 	}
 	i := r.waitingIndex(a.Entry)
 	if i < 0 {

internal/domain/room/room.go

@@ -12,8 +12,8 @@ const (
 	startLead           = 200 * time.Millisecond
 	closeDelay          = 5 * time.Minute
 	autoDJQuiet         = 10 * time.Second
-	refillInterval      = time.Minute
-	ephemeralBurst      = 5
+	RefillInterval      = time.Minute
+	EphemeralBurst      = 5
 	persistentBurst     = 30
 	ephemeralQueueLimit = 10
 	chatLimit           = 500