moteva/server/internal/modules/realtime/outbox_publisher.go

package realtime

import (
	"context"
	"encoding/json"
	"strings"
	"sync"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/zeromicro/go-zero/core/logx"
)

const (
	defaultOutboxInterval = time.Second
	defaultOutboxBatch    = 100
)

type OutboxOptions struct {
	DataSource     string
	IntervalMillis int
	BatchSize      int
}

type OutboxPublisher struct {
	pool     *pgxpool.Pool
	bus      Bus
	interval time.Duration
	batch    int
	ctx      context.Context
	cancel   context.CancelFunc
	wg       sync.WaitGroup
}

func NewOutboxPublisher(ctx context.Context, opts OutboxOptions, bus Bus) (*OutboxPublisher, error) {
	dataSource := strings.TrimSpace(opts.DataSource)
	if dataSource == "" || bus == nil {
		return nil, nil
	}
	pool, err := pgxpool.New(ctx, dataSource)
	if err != nil {
		return nil, err
	}
	interval := time.Duration(opts.IntervalMillis) * time.Millisecond
	if interval <= 0 {
		interval = defaultOutboxInterval
	}
	batch := opts.BatchSize
	if batch <= 0 {
		batch = defaultOutboxBatch
	}
	runCtx, cancel := context.WithCancel(context.Background())
	return &OutboxPublisher{
		pool:     pool,
		bus:      bus,
		interval: interval,
		batch:    batch,
		ctx:      runCtx,
		cancel:   cancel,
	}, nil
}

func (p *OutboxPublisher) Start() error {
	if p == nil || p.pool == nil {
		return nil
	}
	p.wg.Add(1)
	go func() {
		defer p.wg.Done()
		p.loop()
	}()
	return nil
}

func (p *OutboxPublisher) Run() error {
	if p == nil || p.pool == nil {
		return nil
	}
	p.loop()
	return nil
}

func (p *OutboxPublisher) Shutdown() {
	if p == nil {
		return
	}
	if p.cancel != nil {
		p.cancel()
	}
	p.wg.Wait()
	if p.pool != nil {
		p.pool.Close()
	}
}

func (p *OutboxPublisher) loop() {
	ticker := time.NewTicker(p.interval)
	defer ticker.Stop()
	for {
		if err := p.publishPending(p.ctx); err != nil && p.ctx.Err() == nil {
			logx.Errorf("publish outbox events failed: %v", err)
		}
		select {
		case <-p.ctx.Done():
			return
		case <-ticker.C:
		}
	}
}

func (p *OutboxPublisher) publishPending(ctx context.Context) error {
	tx, err := p.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback(ctx)

	rows, err := tx.Query(ctx, `
SELECT id, user_id, aggregate_id, event_type, thread_id, payload_json, created_at
FROM outbox_events
WHERE status = 'pending'
ORDER BY created_at ASC
LIMIT $1
FOR UPDATE SKIP LOCKED
`, p.batch)
	if err != nil {
		return err
	}
	events, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (Event, error) {
		var event Event
		var payload []byte
		if err := row.Scan(&event.ID, &event.UserID, &event.ProjectID, &event.Type, &event.ThreadID, &payload, &event.CreatedAt); err != nil {
			return Event{}, err
		}
		event.Payload = json.RawMessage(payload)
		return event, nil
	})
	if err != nil {
		return err
	}
	for _, event := range events {
		if err := p.bus.Publish(ctx, event); err != nil {
			return err
		}
		if _, err := tx.Exec(ctx, `
UPDATE outbox_events
SET status = 'published', published_at = NOW()
WHERE id = $1
`, event.ID); err != nil {
			return err
		}
	}
	return tx.Commit(ctx)
}