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

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