package dbQueue import ( "context" "f1-game/internal/extend/maps" "f1-game/internal/pb" "runtime/debug" "sync" "sync/atomic" "time" "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/mongo" mopts "go.mongodb.org/mongo-driver/v2/mongo/options" ctime "github.com/cherry-game/cherry/extend/time" dbFacade "f1-game/internal/component/db/facade" cstring "github.com/cherry-game/cherry/extend/string" clog "github.com/cherry-game/cherry/logger" ) type ( TableQueue struct { workers map[string]*tableWorker // worker map db *mongo.Database // gorm db onStopScanSecond time.Duration // 停机时扫描表等待时间 producerWg sync.WaitGroup // 追踪 run (生产者) isRun bool // 是否启动 } tableWorker struct { options tableName string // 表名 dataMaps *maps.OrderedMap[any, *mongo.UpdateOneModel] // 所有待保存的数据 key:hash, value:*Table submitCount int64 // 提交到保存队列的总数 mongoUpdateOpts *mopts.FindOneAndUpdateOptionsBuilder // mongodb update options } options struct { db *mongo.Database // mongodb saveChanSize int // chan的大小 saveTicker time.Duration // 默认保存频率(秒) saveCount int // 默认每次保存的数量 onStopSaveTicker time.Duration // 停机时保存频率(秒) onStopSaveCount int // 停机时每次保存的数量 printSaveLog bool // 是否打印保存日志 } OptionFunc func(options *options) ) func newOption(db *mongo.Database) options { return options{ db: db, saveChanSize: 10240, saveTicker: 60 * time.Second, saveCount: 500, onStopSaveTicker: 1 * time.Second, onStopSaveCount: 1000, printSaveLog: false, } } func NewTableQueue(db *mongo.Database) *TableQueue { if db == nil { clog.Fatal("db is nil") return nil } queue := &TableQueue{ db: db, workers: make(map[string]*tableWorker), onStopScanSecond: 3 * time.Second, isRun: false, } return queue } func (p *TableQueue) AddTable(table dbFacade.Table, optsFunc ...OptionFunc) { if p.isRun { clog.Warn("TableQueue already started") return } if table == nil { clog.Warn("table is nil") return } if cstring.IsBlank(table.TableName()) { clog.Warn("table name is empty.") return } if _, found := p.workers[table.TableName()]; found { clog.Warnf("duplicate table = %s name.", table.TableName()) return } opts := newOption(p.db) for _, optionFunc := range optsFunc { optionFunc(&opts) } worker := &tableWorker{ options: opts, tableName: table.TableName(), dataMaps: maps.NewOrderMap[any, *mongo.UpdateOneModel](opts.saveChanSize, 10), submitCount: 0, mongoUpdateOpts: mopts.FindOneAndUpdate().SetUpsert(true), // set upsert } p.workers[worker.tableName] = worker clog.Infof("add [table = %s, ticker = %d(sec), count = %d]", worker.tableName, int(worker.saveTicker.Seconds()), worker.saveCount, ) } func (p *TableQueue) AddTables(tables []dbFacade.Table, optsFunc ...OptionFunc) { for _, table := range tables { p.AddTable(table, optsFunc...) } } func (p *TableQueue) Start() { if p.isRun { clog.Warn("TableQueue already started") return } p.isRun = true clog.Infof("db queue starting...") for _, table := range p.workers { p.producerWg.Add(1) go table.run(&p.producerWg) } } func (p *TableQueue) Stop() { clog.Infof("stopping db queue....") for _, tw := range p.workers { tw.saveTicker = tw.onStopSaveTicker tw.saveCount = tw.onStopSaveCount } // 阶段 1: 等待所有内存数据刷入 for { if m := p.allWorkerDataCount(); m < 1 { break } // wait... clog.Debugf("saving data! waiting in %d seconds.", int(p.onStopScanSecond.Seconds())) time.Sleep(p.onStopScanSecond) } clog.Infof("stopping db queue finish!") // 阶段 2: 关闭数据源 (这会触发 run 协程退出) for _, worker := range p.workers { worker.dataMaps.Close() } // 阶段 3: 等待所有生产者 (run 协程) 退出 p.producerWg.Wait() } // 对象进入队列 func (p *TableQueue) InQueue(table dbFacade.Table) { worker, found := p.workers[table.TableName()] if !found { clog.Errorf("[table = %s] unregister!", table.TableName()) return } // flush update time table.FlushUpdateTime(ctime.Now().ToMillisecond()) table.CheckField() // marshal to bytes tableBytes, err := bson.Marshal(table) if err != nil { clog.Warnf("[table = %s] Marshal table to bytes error. [err = %v]", table.TableName(), err) return } // set update update := bson.D{{Key: "$set", Value: bson.Raw(tableBytes)}} // set filter filter := bson.M{ "_id": table.PrimaryValue(), // 利用自带的_id主键进行条件过滤 } // build *mongo.UpdateOneModel model := mongo.NewUpdateOneModel().SetFilter(filter).SetUpdate(update).SetUpsert(true) worker.dataMaps.Set(table.PrimaryValue(), model) } func (p *TableQueue) allWorkerDataCount() int { var allDataMapCount = 0 for _, worker := range p.workers { allDataMapCount += worker.dataMaps.Len() } return allDataMapCount } func (p *TableQueue) allSubmitCount() int64 { var allSubmitCount int64 = 0 for _, worker := range p.workers { allSubmitCount += worker.getSubmitCount() } return allSubmitCount } func (p *TableQueue) Report() *pb.DBQueueReportList { list := &pb.DBQueueReportList{} for _, worker := range p.workers { list.List = append(list.List, &pb.DBQueueReport{ TableName: worker.tableName, RemainCount: int64(worker.dataMaps.Len()), SubmitCount: worker.getSubmitCount(), }) } return list } func (p *tableWorker) run(wg *sync.WaitGroup) { defer wg.Done() clog.Infof("run queue for [table = %s]", p.tableName) checkTicker := time.NewTicker(100 * time.Millisecond) defer checkTicker.Stop() lastTicker := ctime.Now() for { select { case <-p.dataMaps.GetStopChan(): return case t := <-checkTicker.C: if t.Unix() >= lastTicker.Add(p.saveTicker).Unix() { lastTicker = ctime.Now() // Use PopBatch to get multiple items at once list := p.dataMaps.PopBatch(p.saveCount) listLen := len(list) if listLen > 0 { p.submitIncrease(listLen) p.bulkWrite(list) p.submitDecrease(listLen) } } } } } func (p *tableWorker) bulkWrite(list []*mongo.UpdateOneModel) { defer func() { if r := recover(); r != nil { clog.Warnf("[table = %s] recover in executor. %s", p.tableName, string(debug.Stack())) for _, model := range list { clog.Warnf("%+v", model) } } }() // 构建批量更新对象 models := make([]mongo.WriteModel, 0, len(list)) for _, model := range list { models = append(models, model) } bulkWriteOptions := mopts.BulkWrite().SetOrdered(false) _, err := p.db.Collection(p.tableName).BulkWrite(context.TODO(), models, bulkWriteOptions) if err != nil && err != mongo.ErrNoDocuments { // ignore error: no document clog.Warnf("[table = %s] save error. [error = %v]", p.tableName, err) for _, model := range list { clog.Warnf("%+v", model) } } else { clog.Debugf("[table = %s, size = %d] queue bulk write.", p.tableName, len(list)) } } func (p *tableWorker) submitIncrease(num int) { atomic.AddInt64(&p.submitCount, int64(num)) } func (p *tableWorker) submitDecrease(num int) { atomic.AddInt64(&p.submitCount, -int64(num)) } func (p *tableWorker) getSubmitCount() int64 { return atomic.LoadInt64(&p.submitCount) } func SetChanSize(saveChanSize int) OptionFunc { return func(options *options) { if saveChanSize > 0 { options.saveChanSize = saveChanSize } } } func SetTicker(saveTicker time.Duration, saveCount int) OptionFunc { return func(options *options) { if saveTicker > 10*time.Millisecond { options.saveTicker = saveTicker } if saveCount > 0 { options.saveCount = saveCount } } } func SetOnStopTicker(onStopSaveTicker time.Duration, onStopSaveCount int) OptionFunc { return func(options *options) { if onStopSaveTicker > 10*time.Millisecond { options.onStopSaveTicker = onStopSaveTicker } if onStopSaveCount > 0 { options.onStopSaveCount = onStopSaveCount } } } func SetPrintSaveLog(printSaveLog bool) OptionFunc { return func(options *options) { options.printSaveLog = printSaveLog } }