mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
fix: cleanup consumed scheduler outbox rows
This commit is contained in:
@@ -8,6 +8,7 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
@@ -16,6 +17,12 @@ type schedulerOutboxRepository struct {
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
type schedulerOutboxCleanupLease struct {
|
||||
conn *sql.Conn
|
||||
}
|
||||
|
||||
const schedulerOutboxDefaultCleanSize = 5000
|
||||
|
||||
func NewSchedulerOutboxRepository(db *sql.DB) service.SchedulerOutboxRepository {
|
||||
return &schedulerOutboxRepository{db: db}
|
||||
}
|
||||
@@ -94,6 +101,60 @@ func (r *schedulerOutboxRepository) MaxID(ctx context.Context) (int64, error) {
|
||||
return maxID, nil
|
||||
}
|
||||
|
||||
func (r *schedulerOutboxRepository) DeleteConsumedUpTo(ctx context.Context, watermark int64, limit int) (int64, error) {
|
||||
if watermark <= 0 {
|
||||
return 0, nil
|
||||
}
|
||||
if limit <= 0 {
|
||||
limit = schedulerOutboxDefaultCleanSize
|
||||
}
|
||||
result, err := r.db.ExecContext(ctx, `
|
||||
WITH doomed AS (
|
||||
SELECT id
|
||||
FROM scheduler_outbox
|
||||
WHERE id <= $1
|
||||
ORDER BY id ASC
|
||||
LIMIT $2
|
||||
)
|
||||
DELETE FROM scheduler_outbox o
|
||||
USING doomed d
|
||||
WHERE o.id = d.id
|
||||
`, watermark, limit)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return result.RowsAffected()
|
||||
}
|
||||
|
||||
func (r *schedulerOutboxRepository) TryAcquireCleanupLock(ctx context.Context) (service.SchedulerOutboxCleanupLease, bool, error) {
|
||||
conn, err := r.db.Conn(ctx)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
|
||||
var acquired bool
|
||||
if err := conn.QueryRowContext(ctx, "SELECT pg_try_advisory_lock(hashtext('scheduler_outbox_cleanup'))").Scan(&acquired); err != nil {
|
||||
_ = conn.Close()
|
||||
return nil, false, err
|
||||
}
|
||||
if !acquired {
|
||||
_ = conn.Close()
|
||||
return nil, false, nil
|
||||
}
|
||||
return &schedulerOutboxCleanupLease{conn: conn}, true, nil
|
||||
}
|
||||
|
||||
func (l *schedulerOutboxCleanupLease) Release() {
|
||||
if l == nil || l.conn == nil {
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
_, _ = l.conn.ExecContext(ctx, "SELECT pg_advisory_unlock(hashtext('scheduler_outbox_cleanup'))")
|
||||
_ = l.conn.Close()
|
||||
l.conn = nil
|
||||
}
|
||||
|
||||
func enqueueSchedulerOutbox(ctx context.Context, exec sqlExecutor, eventType string, accountID *int64, groupID *int64, payload any) error {
|
||||
if exec == nil {
|
||||
return nil
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"regexp"
|
||||
"testing"
|
||||
|
||||
sqlmock "github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSchedulerOutboxRepositoryDeleteConsumedUpToUsesBoundedCTE(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
|
||||
repo := &schedulerOutboxRepository{db: db}
|
||||
const expectedSQL = `
|
||||
WITH doomed AS (
|
||||
SELECT id
|
||||
FROM scheduler_outbox
|
||||
WHERE id <= $1
|
||||
ORDER BY id ASC
|
||||
LIMIT $2
|
||||
)
|
||||
DELETE FROM scheduler_outbox o
|
||||
USING doomed d
|
||||
WHERE o.id = d.id
|
||||
`
|
||||
mock.ExpectExec(regexp.QuoteMeta(expectedSQL)).
|
||||
WithArgs(int64(42), 5000).
|
||||
WillReturnResult(sqlmock.NewResult(0, 17))
|
||||
|
||||
deleted, err := repo.DeleteConsumedUpTo(context.Background(), 42, 5000)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 17, deleted)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestSchedulerOutboxRepositoryDeleteConsumedUpToSkipsNonPositiveWatermark(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
|
||||
repo := &schedulerOutboxRepository{db: db}
|
||||
|
||||
deleted, err := repo.DeleteConsumedUpTo(context.Background(), 0, 5000)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 0, deleted)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestSchedulerOutboxRepositoryTryAcquireCleanupLock(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
|
||||
repo := &schedulerOutboxRepository{db: db}
|
||||
mock.ExpectQuery(regexp.QuoteMeta("SELECT pg_try_advisory_lock(hashtext('scheduler_outbox_cleanup'))")).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"pg_try_advisory_lock"}).AddRow(true))
|
||||
mock.ExpectExec(regexp.QuoteMeta("SELECT pg_advisory_unlock(hashtext('scheduler_outbox_cleanup'))")).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
lease, acquired, err := repo.TryAcquireCleanupLock(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.True(t, acquired)
|
||||
require.NotNil(t, lease)
|
||||
|
||||
lease.Release()
|
||||
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestSchedulerOutboxRepositoryTryAcquireCleanupLockUnavailable(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
|
||||
repo := &schedulerOutboxRepository{db: db}
|
||||
mock.ExpectQuery(regexp.QuoteMeta("SELECT pg_try_advisory_lock(hashtext('scheduler_outbox_cleanup'))")).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"pg_try_advisory_lock"}).AddRow(false))
|
||||
|
||||
lease, acquired, err := repo.TryAcquireCleanupLock(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.False(t, acquired)
|
||||
require.Nil(t, lease)
|
||||
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
@@ -18,4 +18,12 @@ type SchedulerOutboxEvent struct {
|
||||
type SchedulerOutboxRepository interface {
|
||||
ListAfterAndReleaseDedup(ctx context.Context, afterID int64, limit int) ([]SchedulerOutboxEvent, error)
|
||||
MaxID(ctx context.Context) (int64, error)
|
||||
DeleteConsumedUpTo(ctx context.Context, watermark int64, limit int) (int64, error)
|
||||
TryAcquireCleanupLock(ctx context.Context) (SchedulerOutboxCleanupLease, bool, error)
|
||||
}
|
||||
|
||||
// SchedulerOutboxCleanupLease holds the PostgreSQL advisory lock used by
|
||||
// scheduler outbox cleanup.
|
||||
type SchedulerOutboxCleanupLease interface {
|
||||
Release()
|
||||
}
|
||||
|
||||
@@ -0,0 +1,269 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"reflect"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type outboxCleanupCache struct {
|
||||
watermark int64
|
||||
setWatermarks []int64
|
||||
updateErr error
|
||||
}
|
||||
|
||||
func (c *outboxCleanupCache) GetSnapshot(ctx context.Context, bucket SchedulerBucket) ([]*Account, bool, error) {
|
||||
return nil, false, nil
|
||||
}
|
||||
|
||||
func (c *outboxCleanupCache) SetSnapshot(ctx context.Context, bucket SchedulerBucket, accounts []Account) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *outboxCleanupCache) GetAccount(ctx context.Context, accountID int64) (*Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (c *outboxCleanupCache) SetAccount(ctx context.Context, account *Account) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *outboxCleanupCache) DeleteAccount(ctx context.Context, accountID int64) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *outboxCleanupCache) UpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error {
|
||||
return c.updateErr
|
||||
}
|
||||
|
||||
func (c *outboxCleanupCache) TryLockBucket(ctx context.Context, bucket SchedulerBucket, ttl time.Duration) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (c *outboxCleanupCache) UnlockBucket(ctx context.Context, bucket SchedulerBucket) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *outboxCleanupCache) ListBuckets(ctx context.Context) ([]SchedulerBucket, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (c *outboxCleanupCache) GetOutboxWatermark(ctx context.Context) (int64, error) {
|
||||
return c.watermark, nil
|
||||
}
|
||||
|
||||
func (c *outboxCleanupCache) SetOutboxWatermark(ctx context.Context, id int64) error {
|
||||
c.watermark = id
|
||||
c.setWatermarks = append(c.setWatermarks, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
type outboxCleanupDeleteCall struct {
|
||||
watermark int64
|
||||
limit int
|
||||
}
|
||||
|
||||
type outboxCleanupRepo struct {
|
||||
events []SchedulerOutboxEvent
|
||||
rows []int64
|
||||
lockAcquired bool
|
||||
lockAttempts int
|
||||
releaseCount int
|
||||
deleteCalls []outboxCleanupDeleteCall
|
||||
}
|
||||
|
||||
func (r *outboxCleanupRepo) ListAfterAndReleaseDedup(ctx context.Context, afterID int64, limit int) ([]SchedulerOutboxEvent, error) {
|
||||
events := make([]SchedulerOutboxEvent, 0, len(r.events))
|
||||
for _, event := range r.events {
|
||||
if event.ID <= afterID {
|
||||
continue
|
||||
}
|
||||
events = append(events, event)
|
||||
if limit > 0 && len(events) >= limit {
|
||||
break
|
||||
}
|
||||
}
|
||||
return events, nil
|
||||
}
|
||||
|
||||
func (r *outboxCleanupRepo) MaxID(ctx context.Context) (int64, error) {
|
||||
var maxID int64
|
||||
for _, id := range r.rows {
|
||||
if id > maxID {
|
||||
maxID = id
|
||||
}
|
||||
}
|
||||
return maxID, nil
|
||||
}
|
||||
|
||||
func (r *outboxCleanupRepo) DeleteConsumedUpTo(ctx context.Context, watermark int64, limit int) (int64, error) {
|
||||
r.deleteCalls = append(r.deleteCalls, outboxCleanupDeleteCall{
|
||||
watermark: watermark,
|
||||
limit: limit,
|
||||
})
|
||||
if watermark <= 0 || limit <= 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
deleted := int64(0)
|
||||
kept := make([]int64, 0, len(r.rows))
|
||||
for _, id := range r.rows {
|
||||
if id <= watermark && deleted < int64(limit) {
|
||||
deleted++
|
||||
continue
|
||||
}
|
||||
kept = append(kept, id)
|
||||
}
|
||||
r.rows = kept
|
||||
return deleted, nil
|
||||
}
|
||||
|
||||
func (r *outboxCleanupRepo) TryAcquireCleanupLock(ctx context.Context) (SchedulerOutboxCleanupLease, bool, error) {
|
||||
r.lockAttempts++
|
||||
if !r.lockAcquired {
|
||||
return nil, false, nil
|
||||
}
|
||||
return outboxCleanupLease{release: func() {
|
||||
r.releaseCount++
|
||||
}}, true, nil
|
||||
}
|
||||
|
||||
type outboxCleanupLease struct {
|
||||
release func()
|
||||
}
|
||||
|
||||
func (l outboxCleanupLease) Release() {
|
||||
if l.release != nil {
|
||||
l.release()
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerSnapshotServicePollOutboxCleansConsumedRowsAfterWatermark(t *testing.T) {
|
||||
cache := &outboxCleanupCache{}
|
||||
repo := &outboxCleanupRepo{
|
||||
events: []SchedulerOutboxEvent{
|
||||
{ID: 10000, EventType: SchedulerOutboxEventAccountLastUsed},
|
||||
},
|
||||
rows: int64Range(1, 10003),
|
||||
lockAcquired: true,
|
||||
}
|
||||
svc := NewSchedulerSnapshotService(cache, repo, nil, nil, nil)
|
||||
|
||||
svc.pollOutbox()
|
||||
|
||||
if cache.watermark != 10000 {
|
||||
t.Fatalf("expected watermark 10000, got %d", cache.watermark)
|
||||
}
|
||||
if !reflect.DeepEqual(cache.setWatermarks, []int64{10000}) {
|
||||
t.Fatalf("unexpected watermark writes: %#v", cache.setWatermarks)
|
||||
}
|
||||
if !reflect.DeepEqual(repo.rows, []int64{10001, 10002, 10003}) {
|
||||
t.Fatalf("expected rows above watermark to remain, got %#v", repo.rows)
|
||||
}
|
||||
if repo.lockAttempts != 1 || repo.releaseCount != 1 {
|
||||
t.Fatalf("expected one lock acquire/release, got acquire=%d release=%d", repo.lockAttempts, repo.releaseCount)
|
||||
}
|
||||
if len(repo.deleteCalls) != 3 {
|
||||
t.Fatalf("expected cleanup to loop until a short batch, got %d calls", len(repo.deleteCalls))
|
||||
}
|
||||
for _, call := range repo.deleteCalls {
|
||||
if call.watermark != 10000 || call.limit != schedulerOutboxCleanupBatch {
|
||||
t.Fatalf("unexpected cleanup call: %#v", call)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerSnapshotServicePollOutboxSkipsCleanupWhenLockUnavailable(t *testing.T) {
|
||||
cache := &outboxCleanupCache{}
|
||||
repo := &outboxCleanupRepo{
|
||||
events: []SchedulerOutboxEvent{
|
||||
{ID: 3, EventType: SchedulerOutboxEventAccountLastUsed},
|
||||
},
|
||||
rows: []int64{1, 2, 3, 4},
|
||||
lockAcquired: false,
|
||||
}
|
||||
svc := NewSchedulerSnapshotService(cache, repo, nil, nil, nil)
|
||||
|
||||
svc.pollOutbox()
|
||||
|
||||
if cache.watermark != 3 {
|
||||
t.Fatalf("expected watermark 3, got %d", cache.watermark)
|
||||
}
|
||||
if !reflect.DeepEqual(repo.rows, []int64{1, 2, 3, 4}) {
|
||||
t.Fatalf("expected cleanup to skip all rows, got %#v", repo.rows)
|
||||
}
|
||||
if repo.lockAttempts != 1 {
|
||||
t.Fatalf("expected one lock attempt, got %d", repo.lockAttempts)
|
||||
}
|
||||
if len(repo.deleteCalls) != 0 {
|
||||
t.Fatalf("expected no delete calls, got %#v", repo.deleteCalls)
|
||||
}
|
||||
if repo.releaseCount != 0 {
|
||||
t.Fatalf("expected no release without lock, got %d", repo.releaseCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerSnapshotServicePollOutboxDoesNotCleanupOnHandleFailure(t *testing.T) {
|
||||
cache := &outboxCleanupCache{
|
||||
updateErr: errors.New("cache update failed"),
|
||||
}
|
||||
repo := &outboxCleanupRepo{
|
||||
events: []SchedulerOutboxEvent{
|
||||
{
|
||||
ID: 5,
|
||||
EventType: SchedulerOutboxEventAccountLastUsed,
|
||||
Payload: map[string]any{
|
||||
"last_used": map[string]any{"101": float64(123)},
|
||||
},
|
||||
},
|
||||
},
|
||||
rows: []int64{1, 2, 3, 4, 5, 6},
|
||||
lockAcquired: true,
|
||||
}
|
||||
svc := NewSchedulerSnapshotService(cache, repo, nil, nil, nil)
|
||||
|
||||
svc.pollOutbox()
|
||||
|
||||
if len(cache.setWatermarks) != 0 {
|
||||
t.Fatalf("expected no watermark write on handle failure, got %#v", cache.setWatermarks)
|
||||
}
|
||||
if repo.lockAttempts != 0 {
|
||||
t.Fatalf("expected cleanup lock not to be attempted, got %d", repo.lockAttempts)
|
||||
}
|
||||
if len(repo.deleteCalls) != 0 {
|
||||
t.Fatalf("expected no delete calls, got %#v", repo.deleteCalls)
|
||||
}
|
||||
if !reflect.DeepEqual(repo.rows, []int64{1, 2, 3, 4, 5, 6}) {
|
||||
t.Fatalf("expected rows unchanged, got %#v", repo.rows)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerSnapshotServiceCleanupSkipsNonPositiveWatermark(t *testing.T) {
|
||||
repo := &outboxCleanupRepo{
|
||||
rows: []int64{1, 2, 3},
|
||||
lockAcquired: true,
|
||||
}
|
||||
svc := NewSchedulerSnapshotService(&outboxCleanupCache{}, repo, nil, nil, nil)
|
||||
|
||||
svc.cleanupConsumedOutbox(0)
|
||||
|
||||
if repo.lockAttempts != 0 {
|
||||
t.Fatalf("expected no lock attempt for non-positive watermark, got %d", repo.lockAttempts)
|
||||
}
|
||||
if len(repo.deleteCalls) != 0 {
|
||||
t.Fatalf("expected no delete calls, got %#v", repo.deleteCalls)
|
||||
}
|
||||
if !reflect.DeepEqual(repo.rows, []int64{1, 2, 3}) {
|
||||
t.Fatalf("expected rows unchanged, got %#v", repo.rows)
|
||||
}
|
||||
}
|
||||
|
||||
func int64Range(start, end int64) []int64 {
|
||||
values := make([]int64, 0, end-start+1)
|
||||
for id := start; id <= end; id++ {
|
||||
values = append(values, id)
|
||||
}
|
||||
return values
|
||||
}
|
||||
@@ -18,7 +18,10 @@ var (
|
||||
ErrSchedulerFallbackLimited = errors.New("scheduler db fallback limited")
|
||||
)
|
||||
|
||||
const outboxEventTimeout = 2 * time.Minute
|
||||
const (
|
||||
outboxEventTimeout = 2 * time.Minute
|
||||
schedulerOutboxCleanupBatch = 5000
|
||||
)
|
||||
|
||||
// batchSeenKey tracks which (groupID, platform) bucket sets have already been
|
||||
// rebuilt within a single pollOutbox call, to avoid redundant work when multiple
|
||||
@@ -280,11 +283,42 @@ func (s *SchedulerSnapshotService) pollOutbox() {
|
||||
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox watermark write failed: %v", wmErr)
|
||||
} else {
|
||||
watermarkForCheck = lastID
|
||||
s.cleanupConsumedOutbox(lastID)
|
||||
}
|
||||
|
||||
s.checkOutboxLag(ctx, events[0], watermarkForCheck)
|
||||
}
|
||||
|
||||
func (s *SchedulerSnapshotService) cleanupConsumedOutbox(watermark int64) {
|
||||
if s == nil || s.outboxRepo == nil || watermark <= 0 {
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
lease, acquired, err := s.outboxRepo.TryAcquireCleanupLock(ctx)
|
||||
if err != nil {
|
||||
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox cleanup lock failed: %v", err)
|
||||
return
|
||||
}
|
||||
if !acquired {
|
||||
return
|
||||
}
|
||||
defer lease.Release()
|
||||
|
||||
for {
|
||||
deleted, err := s.outboxRepo.DeleteConsumedUpTo(ctx, watermark, schedulerOutboxCleanupBatch)
|
||||
if err != nil {
|
||||
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox cleanup failed: watermark=%d err=%v", watermark, err)
|
||||
return
|
||||
}
|
||||
if deleted == 0 || deleted < schedulerOutboxCleanupBatch {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *SchedulerSnapshotService) handleOutboxEvent(ctx context.Context, event SchedulerOutboxEvent, seen map[batchSeenKey]struct{}) error {
|
||||
switch event.EventType {
|
||||
case SchedulerOutboxEventAccountLastUsed:
|
||||
|
||||
Reference in New Issue
Block a user