github.com/lingyao2333/mo-zero@v1.4.1/core/stores/mon/bulkinserter.go (about)

     1  package mon
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/lingyao2333/mo-zero/core/executors"
     8  	"github.com/lingyao2333/mo-zero/core/logx"
     9  	"go.mongodb.org/mongo-driver/mongo"
    10  )
    11  
    12  const (
    13  	flushInterval = time.Second
    14  	maxBulkRows   = 1000
    15  )
    16  
    17  type (
    18  	// ResultHandler is a handler that used to handle results.
    19  	ResultHandler func(*mongo.InsertManyResult, error)
    20  
    21  	// A BulkInserter is used to insert bulk of mongo records.
    22  	BulkInserter struct {
    23  		executor *executors.PeriodicalExecutor
    24  		inserter *dbInserter
    25  	}
    26  )
    27  
    28  // NewBulkInserter returns a BulkInserter.
    29  func NewBulkInserter(coll *mongo.Collection, interval ...time.Duration) *BulkInserter {
    30  	inserter := &dbInserter{
    31  		collection: coll,
    32  	}
    33  
    34  	duration := flushInterval
    35  	if len(interval) > 0 {
    36  		duration = interval[0]
    37  	}
    38  
    39  	return &BulkInserter{
    40  		executor: executors.NewPeriodicalExecutor(duration, inserter),
    41  		inserter: inserter,
    42  	}
    43  }
    44  
    45  // Flush flushes the inserter, writes all pending records.
    46  func (bi *BulkInserter) Flush() {
    47  	bi.executor.Flush()
    48  }
    49  
    50  // Insert inserts doc.
    51  func (bi *BulkInserter) Insert(doc interface{}) {
    52  	bi.executor.Add(doc)
    53  }
    54  
    55  // SetResultHandler sets the result handler.
    56  func (bi *BulkInserter) SetResultHandler(handler ResultHandler) {
    57  	bi.executor.Sync(func() {
    58  		bi.inserter.resultHandler = handler
    59  	})
    60  }
    61  
    62  type dbInserter struct {
    63  	collection    *mongo.Collection
    64  	documents     []interface{}
    65  	resultHandler ResultHandler
    66  }
    67  
    68  func (in *dbInserter) AddTask(doc interface{}) bool {
    69  	in.documents = append(in.documents, doc)
    70  	return len(in.documents) >= maxBulkRows
    71  }
    72  
    73  func (in *dbInserter) Execute(objs interface{}) {
    74  	docs := objs.([]interface{})
    75  	if len(docs) == 0 {
    76  		return
    77  	}
    78  
    79  	result, err := in.collection.InsertMany(context.Background(), docs)
    80  	if in.resultHandler != nil {
    81  		in.resultHandler(result, err)
    82  	} else if err != nil {
    83  		logx.Error(err)
    84  	}
    85  }
    86  
    87  func (in *dbInserter) RemoveAll() interface{} {
    88  	documents := in.documents
    89  	in.documents = nil
    90  	return documents
    91  }