junhong_cmp_fiber/internal/task/collection_batch_consumer.go

package task

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

	"go.uber.org/zap"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"

	"github.com/break/junhong_cmp_fiber/internal/model"
	"github.com/break/junhong_cmp_fiber/pkg/constants"
)

// collectionWriteTimeout 是一批采集记录写库的超时。
const collectionWriteTimeout = 30 * time.Second

// CollectionBatchConsumer 消费采集记录消息（topic card.collection），攒满一批或等满 1 秒写一次库：
// 写成功后整批确认，写失败整批延迟重投；重复记录由唯一索引忽略。写库只占一个连接。
type CollectionBatchConsumer struct {
	db        *gorm.DB
	batchSize int
	logger    *zap.Logger
	items     chan *collectionItem
}

type collectionItem struct {
	msg  CollectionMessage
	done chan error
}

// NewCollectionBatchConsumer 创建采集记录批量消费者。
func NewCollectionBatchConsumer(db *gorm.DB, batchSize int, logger *zap.Logger) *CollectionBatchConsumer {
	if batchSize <= 0 {
		batchSize = constants.PollingCollectionDefaultBatchSize
	}
	if logger == nil {
		logger = zap.NewNop()
	}
	return &CollectionBatchConsumer{db: db, batchSize: batchSize, logger: logger,
		items: make(chan *collectionItem, constants.PollingCollectionMaxInFlight)}
}

// Start 启动攒批写库循环，ctx 结束时写完已缓冲的一批后退出。
func (c *CollectionBatchConsumer) Start(ctx context.Context) {
	go c.run(ctx)
}

// Handle 处理一条采集记录消息：交给攒批循环并等待本批写库结果。
func (c *CollectionBatchConsumer) Handle(ctx context.Context, body []byte, attempts uint16, _ func()) time.Duration {
	var msg CollectionMessage
	if err := json.Unmarshal(body, &msg); err != nil || (msg.Success == nil && msg.Failure == nil) {
		c.logger.Warn("采集记录消息无效，丢弃", zap.ByteString("body", body))
		return 0
	}
	item := &collectionItem{msg: msg, done: make(chan error, 1)}
	select {
	case c.items <- item:
	case <-ctx.Done():
		return ingestRetryDelay(attempts)
	}
	select {
	case err := <-item.done:
		if err == nil {
			return 0
		}
		delay := ingestRetryDelay(attempts)
		if attempts >= constants.PollingIngestErrorLogFromAttempt {
			c.logger.Error("写采集记录持续失败，继续重投", zap.Uint16("attempts", attempts),
				zap.String("observation_id", collectionObservationID(msg)), zap.Duration("retry_after", delay), zap.Error(err))
		}
		return delay
	case <-ctx.Done():
		return ingestRetryDelay(attempts)
	}
}

func (c *CollectionBatchConsumer) run(ctx context.Context) {
	ticker := time.NewTicker(constants.PollingCollectionFlushInterval)
	defer ticker.Stop()
	batch := make([]*collectionItem, 0, c.batchSize)
	flush := func() {
		if len(batch) == 0 {
			return
		}
		c.flush(batch)
		batch = make([]*collectionItem, 0, c.batchSize)
	}
	for {
		select {
		case <-ctx.Done():
			flush()
			return
		case item := <-c.items:
			batch = append(batch, item)
			if len(batch) >= c.batchSize {
				flush()
			}
		case <-ticker.C:
			flush()
		}
	}
}

// flush 一次写入一批成功记录与失败记录，并把结果通知本批每条消息。
func (c *CollectionBatchConsumer) flush(batch []*collectionItem) {
	var successes []*model.CardTrafficCollection
	var failures []*model.CardTrafficCollectionFailure
	for _, item := range batch {
		if item.msg.Success != nil {
			successes = append(successes, item.msg.Success)
		}
		if item.msg.Failure != nil {
			failures = append(failures, item.msg.Failure)
		}
	}
	ctx, cancel := context.WithTimeout(context.Background(), collectionWriteTimeout)
	defer cancel()
	err := c.write(ctx, successes, failures)
	if err != nil {
		c.logger.Warn("批量写采集记录失败，整批重投", zap.Int("success_rows", len(successes)),
			zap.Int("failure_rows", len(failures)), zap.Error(err))
	}
	for _, item := range batch {
		item.done <- err
	}
}

func (c *CollectionBatchConsumer) write(ctx context.Context, successes []*model.CardTrafficCollection, failures []*model.CardTrafficCollectionFailure) error {
	return c.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if len(successes) > 0 {
			if err := tx.Clauses(clause.OnConflict{DoNothing: true}).CreateInBatches(successes, len(successes)).Error; err != nil {
				return err
			}
		}
		if len(failures) > 0 {
			if err := tx.Clauses(clause.OnConflict{DoNothing: true}).CreateInBatches(failures, len(failures)).Error; err != nil {
				return err
			}
		}
		return nil
	})
}

func collectionObservationID(msg CollectionMessage) string {
	if msg.Success != nil {
		return msg.Success.ObservationID
	}
	if msg.Failure != nil {
		return msg.Failure.ObservationID
	}
	return ""
}