package postgres import ( "context" "errors" "fmt" "sort" "time" "telesrv/internal/domain" "telesrv/internal/store" "telesrv/internal/store/postgres/sqlcgen" ) // defaultDispatchLease 是 'dispatching' 行被判定租约过期、可被重新 claim 的默认时长。 // 与 docs/message-module.md 的 outbox 背压参数对应;生产由 config 注入覆盖。 const ( defaultDispatchLease = 30 * time.Second defaultDispatchPoisonCleanupBatch = 256 maxDispatchPoisonCleanupBatch = 1000 ) var errInvalidDispatchOutboxExclusionPair = errors.New("dispatch outbox exclusion requires both raw auth key and session id") // enqueueDispatch is the only production write boundary for dispatch_outbox. // A zero pair means no originating session is excluded; a non-zero pair identifies // one exact physical raw-auth/session tuple. A half pair is never meaningful because // session IDs are not globally unique and must fail the surrounding transaction. func enqueueDispatch(ctx context.Context, q *sqlcgen.Queries, arg sqlcgen.EnqueueDispatchParams) error { hasAuthKey := arg.ExcludeAuthKeyID != 0 hasSession := arg.ExcludeSessionID != 0 if hasAuthKey != hasSession { return errInvalidDispatchOutboxExclusionPair } return q.EnqueueDispatch(ctx, arg) } // DispatchOutboxStore 用 PostgreSQL 实现 transactional outbox。 type DispatchOutboxStore struct { db sqlcgen.DBTX q *sqlcgen.Queries leaseSeconds int32 } // DispatchOutboxOption 调整 DispatchOutboxStore 的 claim 行为。 type DispatchOutboxOption func(*DispatchOutboxStore) // WithLeaseTimeout 设置租约超时;<=0 时保持默认。 func WithLeaseTimeout(d time.Duration) DispatchOutboxOption { return func(s *DispatchOutboxStore) { if d > 0 { s.leaseSeconds = int32(d / time.Second) if s.leaseSeconds < 1 { s.leaseSeconds = 1 } } } } // NewDispatchOutboxStore 基于 pgx 连接池(或事务)创建 DispatchOutboxStore。 func NewDispatchOutboxStore(db sqlcgen.DBTX, opts ...DispatchOutboxOption) *DispatchOutboxStore { s := &DispatchOutboxStore{ db: db, q: sqlcgen.New(db), leaseSeconds: int32(defaultDispatchLease / time.Second), } for _, opt := range opts { if opt != nil { opt(s) } } return s } func (s *DispatchOutboxStore) ClaimPending(ctx context.Context, limit int) ([]store.DispatchOutboxItem, error) { if limit <= 0 { limit = 100 } if limit > 1000 { limit = 1000 } rows, err := s.q.ClaimDispatchOutbox(ctx, sqlcgen.ClaimDispatchOutboxParams{ LeaseSeconds: s.leaseSeconds, LimitCount: int32(limit), }) if err != nil { return nil, fmt.Errorf("claim dispatch outbox: %w", err) } return dispatchItemsFromClaimRows(rows), nil } // ClaimPendingShards 只领取固定 logical shard 集合中的用户 head 事件。 // shardCount 是稳定哈希空间,shardIDs 是当前 worker 独占的子集;worker 数变化只改变 // shard→worker 的运行时归属,不改变 user→shard,从而避免同一用户被并行领取。 func (s *DispatchOutboxStore) ClaimPendingShards(ctx context.Context, shardCount int, shardIDs []int, limit int) ([]store.DispatchOutboxItem, error) { if shardCount <= 0 || len(shardIDs) == 0 { return nil, nil } if shardCount != store.DispatchOutboxLogicalShards { return nil, fmt.Errorf("claim dispatch outbox shards: shard count %d, want stable %d", shardCount, store.DispatchOutboxLogicalShards) } if limit <= 0 { limit = 100 } if limit > 1000 { limit = 1000 } ids := make([]int16, 0, len(shardIDs)) seen := make(map[int]struct{}, len(shardIDs)) for _, id := range shardIDs { if id < 0 || id >= shardCount { continue } if _, ok := seen[id]; ok { continue } seen[id] = struct{}{} ids = append(ids, int16(id)) } if len(ids) == 0 { return nil, nil } rows, err := s.q.ClaimDispatchOutboxShards(ctx, sqlcgen.ClaimDispatchOutboxShardsParams{ LeaseSeconds: s.leaseSeconds, LimitCount: int32(limit), ShardIds: ids, }) if err != nil { return nil, fmt.Errorf("claim dispatch outbox shards: %w", err) } out := make([]store.DispatchOutboxItem, 0, len(rows)) for _, row := range rows { out = append(out, store.DispatchOutboxItem{ ID: row.ID, TargetUserID: row.TargetUserID, Pts: int(row.Pts), EventType: domain.UpdateEventType(row.EventType), ExcludeAuthKeyID: authKeyIDFromInt64(row.ExcludeAuthKeyID), ExcludeSessionID: row.ExcludeSessionID, Attempts: int(row.Attempts), }) } return out, nil } func dispatchItemsFromClaimRows(rows []sqlcgen.ClaimDispatchOutboxRow) []store.DispatchOutboxItem { out := make([]store.DispatchOutboxItem, 0, len(rows)) for _, row := range rows { out = append(out, store.DispatchOutboxItem{ ID: row.ID, TargetUserID: row.TargetUserID, Pts: int(row.Pts), EventType: domain.UpdateEventType(row.EventType), ExcludeAuthKeyID: authKeyIDFromInt64(row.ExcludeAuthKeyID), ExcludeSessionID: row.ExcludeSessionID, Attempts: int(row.Attempts), }) } return out } // MarkDeliveredBatch 一次性删除一批已投递的 outbox 行(方案 A:投递成功即删),取代逐条 MarkDelivered。 func (s *DispatchOutboxStore) MarkDeliveredBatch(ctx context.Context, items []store.DispatchOutboxItem) error { if len(items) == 0 { return nil } targetUserIDs := make([]int64, len(items)) ids := make([]int64, len(items)) expectedAttempts := make([]int32, len(items)) for i, it := range items { targetUserIDs[i] = it.TargetUserID ids[i] = it.ID expectedAttempts[i] = int32(it.Attempts) } rows, err := s.withExclusiveLaneFences(ctx, targetUserIDs, func(db sqlcgen.DBTX) (int64, error) { return sqlcgen.New(db).MarkDispatchDeliveredBatch(ctx, sqlcgen.MarkDispatchDeliveredBatchParams{ TargetUserIds: targetUserIDs, Ids: ids, ExpectedAttempts: expectedAttempts, }) }) if err != nil { return fmt.Errorf("mark dispatch delivered batch: %w", err) } if rows != int64(len(items)) { return fmt.Errorf("mark dispatch delivered batch: %w: updated %d of %d", store.ErrDispatchLeaseLost, rows, len(items)) } return nil } func (s *DispatchOutboxStore) MarkDelivered(ctx context.Context, item store.DispatchOutboxItem) error { rows, err := s.withExclusiveLaneFences(ctx, []int64{item.TargetUserID}, func(db sqlcgen.DBTX) (int64, error) { return sqlcgen.New(db).MarkDispatchDelivered(ctx, sqlcgen.MarkDispatchDeliveredParams{ TargetUserID: item.TargetUserID, ID: item.ID, ExpectedAttempts: int32(item.Attempts), }) }) if err != nil { return fmt.Errorf("mark dispatch delivered: %w", err) } if rows != 1 { return fmt.Errorf("mark dispatch delivered: %w", store.ErrDispatchLeaseLost) } return nil } // withExclusiveLaneFences serializes the empty-lane transition with producers' // shared append fences. The DELETE runs as a later READ COMMITTED statement, so // it sees every producer that committed before the exclusive fence was granted. func (s *DispatchOutboxStore) withExclusiveLaneFences( ctx context.Context, userIDs []int64, work func(sqlcgen.DBTX) (int64, error), ) (int64, error) { beginner, ok := s.db.(txBeginner) if !ok { return 0, fmt.Errorf("dispatch lane transition requires transaction-capable database") } tx, err := beginner.Begin(ctx) if err != nil { return 0, fmt.Errorf("begin dispatch lane transition: %w", err) } committed := false defer func() { if !committed { _ = tx.Rollback(ctx) } }() if err := lockDispatchOutboxLanesExclusive(ctx, tx, userIDs); err != nil { return 0, err } rows, err := work(tx) if err != nil { return 0, err } if err := tx.Commit(ctx); err != nil { return 0, fmt.Errorf("commit dispatch lane transition: %w", err) } committed = true return rows, nil } func lockDispatchOutboxLanesExclusive(ctx context.Context, db sqlcgen.DBTX, userIDs []int64) error { unique := make([]int64, 0, len(userIDs)) seen := make(map[int64]struct{}, len(userIDs)) for _, userID := range userIDs { if userID <= 0 { continue } if _, ok := seen[userID]; ok { continue } seen[userID] = struct{}{} unique = append(unique, userID) } if len(unique) == 0 { return nil } sort.Slice(unique, func(i, j int) bool { return unique[i] < unique[j] }) if _, err := db.Exec(ctx, ` SELECT pg_advisory_xact_lock(dispatch_outbox_lane_advisory_key(streams.target_user_id)) FROM unnest($1::bigint[]) AS streams(target_user_id) ORDER BY streams.target_user_id`, unique); err != nil { return fmt.Errorf("lock dispatch outbox lane transition: %w", err) } return nil } func (s *DispatchOutboxStore) MarkFailed(ctx context.Context, item store.DispatchOutboxItem, lastError string) error { rows, err := s.q.MarkDispatchFailed(ctx, sqlcgen.MarkDispatchFailedParams{ TargetUserID: item.TargetUserID, ID: item.ID, LastError: lastError, ExpectedAttempts: int32(item.Attempts), }) if err != nil { return fmt.Errorf("mark dispatch failed: %w", err) } if rows != 1 { return fmt.Errorf("mark dispatch failed: %w", store.ErrDispatchLeaseLost) } return nil } func (s *DispatchOutboxStore) DeleteFailed(ctx context.Context, olderThan time.Duration, limit int) (int, error) { if olderThan <= 0 { olderThan = time.Minute } if limit <= 0 { limit = defaultDispatchPoisonCleanupBatch } if limit > maxDispatchPoisonCleanupBatch { limit = maxDispatchPoisonCleanupBatch } olderThanSeconds := int32(olderThan / time.Second) rows, err := s.db.Query(ctx, ` SELECT h.target_user_id FROM dispatch_outbox_user_heads h WHERE h.status = 'failed' AND h.updated_at < now() - make_interval(secs => $1::int) ORDER BY h.updated_at ASC, h.target_user_id ASC, h.head_id ASC LIMIT $2`, olderThanSeconds, int32(limit)) if err != nil { return 0, fmt.Errorf("list failed dispatch outbox lanes: %w", err) } userIDs := make([]int64, 0, limit) for rows.Next() { var userID int64 if err := rows.Scan(&userID); err != nil { rows.Close() return 0, fmt.Errorf("scan failed dispatch outbox lane: %w", err) } userIDs = append(userIDs, userID) } if err := rows.Err(); err != nil { rows.Close() return 0, fmt.Errorf("iterate failed dispatch outbox lanes: %w", err) } rows.Close() if len(userIDs) == 0 { return 0, nil } deleted, err := s.withExclusiveLaneFences(ctx, userIDs, func(db sqlcgen.DBTX) (int64, error) { count, deleteErr := sqlcgen.New(db).DeleteFailedDispatchOutbox(ctx, sqlcgen.DeleteFailedDispatchOutboxParams{ OlderThanSeconds: olderThanSeconds, LimitCount: int32(limit), TargetUserIds: userIDs, }) return int64(count), deleteErr }) if err != nil { return 0, fmt.Errorf("delete failed dispatch outbox: %w", err) } return int(deleted), nil }