s_cherry_e/sr/internal/component/db/queue/table_queue.go

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"

	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(time.Now().UnixMilli())
	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 := time.Now()

	for {
		select {
		case <-p.dataMaps.GetStopChan():
			return
		case t := <-checkTicker.C:
			if t.Unix() >= lastTicker.Add(p.saveTicker).Unix() {
				lastTicker = time.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
	}
}