batch_image_queue_test.go 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199
  1. //go:build unit
  2. package repository
  3. import (
  4. "context"
  5. "errors"
  6. "testing"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/service"
  9. "github.com/alicebob/miniredis/v2"
  10. "github.com/redis/go-redis/v9"
  11. "github.com/stretchr/testify/require"
  12. )
  13. func TestBatchImageQueue_DuplicateEnqueueReturnsAlreadyQueued(t *testing.T) {
  14. ctx := context.Background()
  15. queue, _ := newBatchImageQueueTest(t)
  16. batchID := "imgbatch_duplicate"
  17. require.NoError(t, queue.Enqueue(ctx, batchID))
  18. err := queue.Enqueue(ctx, batchID)
  19. require.Error(t, err)
  20. require.True(t, errors.Is(err, service.ErrBatchImageAlreadyQueued))
  21. }
  22. func TestBatchImageQueue_RequeueAfterMovesJobFromActiveToDelayed(t *testing.T) {
  23. ctx := context.Background()
  24. queue, _ := newBatchImageQueueTest(t)
  25. batchID := "imgbatch_requeue_after"
  26. require.NoError(t, queue.rdb.ZAdd(ctx, queue.activeKey, redis.Z{
  27. Score: float64(time.Now().UnixMilli()),
  28. Member: batchID,
  29. }).Err())
  30. require.NoError(t, queue.RequeueAfter(ctx, batchID, time.Minute))
  31. require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.activeKey, batchID).Err(), redis.Nil)
  32. score, err := queue.rdb.ZScore(ctx, queue.delayedKey, batchID).Result()
  33. require.NoError(t, err)
  34. require.Greater(t, score, float64(time.Now().UnixMilli()))
  35. }
  36. func TestBatchImageQueue_MoveDueDelayedToReadyMovesDueJobs(t *testing.T) {
  37. ctx := context.Background()
  38. queue, _ := newBatchImageQueueTest(t)
  39. dueBatchID := "imgbatch_due"
  40. futureBatchID := "imgbatch_future"
  41. now := time.Now()
  42. require.NoError(t, queue.rdb.ZAdd(ctx, queue.delayedKey,
  43. redis.Z{Score: float64(now.Add(-time.Second).UnixMilli()), Member: dueBatchID},
  44. redis.Z{Score: float64(now.Add(time.Hour).UnixMilli()), Member: futureBatchID},
  45. ).Err())
  46. moved, err := queue.MoveDueDelayedToReady(ctx, 10)
  47. require.NoError(t, err)
  48. require.Equal(t, 1, moved)
  49. require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.delayedKey, dueBatchID).Err(), redis.Nil)
  50. require.NoError(t, queue.rdb.ZScore(ctx, queue.delayedKey, futureBatchID).Err())
  51. reserved, err := queue.Reserve(ctx, time.Millisecond)
  52. require.NoError(t, err)
  53. require.Equal(t, dueBatchID, reserved.BatchID)
  54. }
  55. func TestBatchImageQueue_RecoverStaleActiveMovesStaleJobsToReady(t *testing.T) {
  56. ctx := context.Background()
  57. queue, _ := newBatchImageQueueTest(t)
  58. staleBatchID := "imgbatch_stale"
  59. recentBatchID := "imgbatch_recent"
  60. now := time.Now()
  61. require.NoError(t, queue.rdb.ZAdd(ctx, queue.activeKey,
  62. redis.Z{Score: float64(now.Add(-time.Hour).UnixMilli()), Member: staleBatchID},
  63. redis.Z{Score: float64(now.UnixMilli()), Member: recentBatchID},
  64. ).Err())
  65. moved, err := queue.RecoverStaleActive(ctx, 10*time.Minute, 10)
  66. require.NoError(t, err)
  67. require.Equal(t, 1, moved)
  68. require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.activeKey, staleBatchID).Err(), redis.Nil)
  69. require.NoError(t, queue.rdb.ZScore(ctx, queue.activeKey, recentBatchID).Err())
  70. reserved, err := queue.Reserve(ctx, time.Millisecond)
  71. require.NoError(t, err)
  72. require.Equal(t, staleBatchID, reserved.BatchID)
  73. }
  74. func TestBatchImageQueue_JobLockReleaseOnlyDeletesMatchingToken(t *testing.T) {
  75. ctx := context.Background()
  76. queue, _ := newBatchImageQueueTest(t)
  77. batchID := "imgbatch_lock"
  78. lock, ok, err := queue.TryAcquireJobLock(ctx, batchID, time.Minute)
  79. require.NoError(t, err)
  80. require.True(t, ok)
  81. require.NoError(t, queue.rdb.Set(ctx, queue.lockKey(batchID), "other-token", time.Minute).Err())
  82. require.NoError(t, lock.Release(ctx))
  83. got, err := queue.rdb.Get(ctx, queue.lockKey(batchID)).Result()
  84. require.NoError(t, err)
  85. require.Equal(t, "other-token", got)
  86. require.NoError(t, queue.rdb.Del(ctx, queue.lockKey(batchID)).Err())
  87. lock, ok, err = queue.TryAcquireJobLock(ctx, batchID, time.Minute)
  88. require.NoError(t, err)
  89. require.True(t, ok)
  90. require.NoError(t, lock.Release(ctx))
  91. require.ErrorIs(t, queue.rdb.Get(ctx, queue.lockKey(batchID)).Err(), redis.Nil)
  92. }
  93. func TestBatchImageQueue_ReserveAtomicallyMovesJobToActive(t *testing.T) {
  94. ctx := context.Background()
  95. queue, _ := newBatchImageQueueTest(t)
  96. batchID := "imgbatch_reserve"
  97. require.NoError(t, queue.Enqueue(ctx, batchID))
  98. reserved, err := queue.Reserve(ctx, time.Second)
  99. require.NoError(t, err)
  100. require.Equal(t, batchID, reserved.BatchID)
  101. // 弹出与写入 active 必须原子完成:ready 已空,active 中有该 job。
  102. require.Equal(t, int64(0), queue.rdb.LLen(ctx, queue.readyKey).Val())
  103. score, err := queue.rdb.ZScore(ctx, queue.activeKey, batchID).Result()
  104. require.NoError(t, err)
  105. require.Positive(t, score)
  106. }
  107. func TestBatchImageQueue_ReserveReturnsEmptyAfterTimeout(t *testing.T) {
  108. ctx := context.Background()
  109. queue, _ := newBatchImageQueueTest(t)
  110. start := time.Now()
  111. _, err := queue.Reserve(ctx, 50*time.Millisecond)
  112. require.ErrorIs(t, err, service.ErrBatchImageQueueEmpty)
  113. require.Less(t, time.Since(start), 5*time.Second)
  114. }
  115. func TestBatchImageQueue_ReserveDropsInvalidPayload(t *testing.T) {
  116. ctx := context.Background()
  117. queue, _ := newBatchImageQueueTest(t)
  118. require.NoError(t, queue.rdb.LPush(ctx, queue.readyKey, "not-a-batch-id").Err())
  119. _, err := queue.Reserve(ctx, 10*time.Millisecond)
  120. require.ErrorIs(t, err, service.ErrInvalidBatchImageQueuePayload)
  121. // 非法 payload 不得残留在 active zset,否则 stale 恢复会无限重投。
  122. require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.activeKey, "not-a-batch-id").Err(), redis.Nil)
  123. }
  124. func TestBatchImageQueue_HeartbeatOnlyRefreshesExistingActiveMember(t *testing.T) {
  125. ctx := context.Background()
  126. queue, _ := newBatchImageQueueTest(t)
  127. batchID := "imgbatch_heartbeat"
  128. // 不在 active 中:心跳不得创建幽灵成员。
  129. require.NoError(t, queue.Heartbeat(ctx, batchID))
  130. require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.activeKey, batchID).Err(), redis.Nil)
  131. require.NoError(t, queue.rdb.ZAdd(ctx, queue.activeKey, redis.Z{Score: 1, Member: batchID}).Err())
  132. require.NoError(t, queue.Heartbeat(ctx, batchID))
  133. score, err := queue.rdb.ZScore(ctx, queue.activeKey, batchID).Result()
  134. require.NoError(t, err)
  135. require.Greater(t, score, float64(1))
  136. }
  137. func TestBatchImageQueue_JobLockRefreshExtendsTTLOnlyForHolder(t *testing.T) {
  138. ctx := context.Background()
  139. queue, mr := newBatchImageQueueTest(t)
  140. batchID := "imgbatch_lock_refresh"
  141. lock, ok, err := queue.TryAcquireJobLock(ctx, batchID, time.Minute)
  142. require.NoError(t, err)
  143. require.True(t, ok)
  144. refresher, isRefresher := lock.(service.BatchImageJobLockRefresher)
  145. require.True(t, isRefresher)
  146. require.NoError(t, refresher.Refresh(ctx, 10*time.Minute))
  147. ttl := mr.TTL(queue.lockKey(batchID))
  148. require.Greater(t, ttl, 5*time.Minute)
  149. // token 不匹配时不得续期他人持有的锁。
  150. require.NoError(t, queue.rdb.Set(ctx, queue.lockKey(batchID), "other-token", time.Minute).Err())
  151. require.NoError(t, refresher.Refresh(ctx, 10*time.Minute))
  152. ttl = mr.TTL(queue.lockKey(batchID))
  153. require.LessOrEqual(t, ttl, time.Minute)
  154. }
  155. func newBatchImageQueueTest(t *testing.T) (*batchImageQueue, *miniredis.Miniredis) {
  156. t.Helper()
  157. mr := miniredis.RunT(t)
  158. rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
  159. t.Cleanup(func() {
  160. _ = rdb.Close()
  161. })
  162. queue := newBatchImageQueueWithOptions(rdb, batchImageQueueOptions{
  163. InflightTTL: time.Hour,
  164. LockTTL: time.Minute,
  165. })
  166. return queue, mr
  167. }