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 "" }