package sender import ( "context" "fmt" "log" "time" m "git.k-space.ee/k-space/logmower-shipper/pkg/mongo" "github.com/jtagcat/util/batch" "github.com/jtagcat/util/retry" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" "k8s.io/apimachinery/pkg/util/wait" ) var ( Simulate = false MaxBatchItems = 1000 ) const ( MaxBatchTime = 5 * time.Second ) // wrapper to force copying before use func backoff() wait.Backoff { return wait.Backoff{ Duration: 2 * time.Second, Factor: 4, Jitter: 0.2, Cap: 2 * time.Minute, } } type Queue <-chan m.Record func (queue Queue) Sender(db *mongo.Collection, metricsFilename string, cancelOnError func()) { batched := make(chan []m.Record) // metrics for batcher and queue go func() { ctx, cancel := context.WithCancel(context.Background()) defer cancel() go func() { for { promShipperQueued.WithLabelValues(metricsFilename).Set(float64( len(queue))) timer := time.NewTimer(time.Second) select { case <-ctx.Done(): return case <-timer.C: } } }() batch.Batch(MaxBatchItems, MaxBatchTime, queue, batched) // returns when sendQueue is closed }() for { promShipperSynced.WithLabelValues(metricsFilename).Set(1) batch, ok := <-batched if !ok { return } promShipperBatchSizeResult.Observe(float64(len(batch))) promShipperSynced.WithLabelValues(metricsFilename).Set(0) err := retry.OnError(backoff(), func() (_ bool, _ error) { result, err := insertManyWithSimulate(db, batch) var succeedCount int if result != nil { succeedCount = len(result.InsertedIDs) } promShipperDbSent.WithLabelValues(metricsFilename).Add(float64(succeedCount)) if err == nil { return } promShipperDbSendError.WithLabelValues(metricsFilename).Add(1) batch = batch[succeedCount:] if succeedCount == len(batch) { return true, fmt.Errorf("all insertions in batch were successful, yet failure in database: %w", err) } return true, fmt.Errorf("insert record with offset %d to database: %w", batch[0].Offset, err) }) if err != nil { log.Printf("batch insert %q to mongo: %e", metricsFilename, err) cancelOnError() return } } } func insertManyWithSimulate(db *mongo.Collection, batch []m.Record) (*mongo.InsertManyResult, error) { if !Simulate { var batchBson []interface{} // mongo does not like typing for _, b := range batch { batchBson = append(batchBson, b.ToBson()) } tru := true return db.InsertMany(m.GlobalTimeout(context.Background()), batchBson, &options.InsertManyOptions{Ordered: &tru}) } fmt.Printf("simulating successful database bulk write: %v", batch) var res mongo.InsertManyResult for range batch { res.InsertedIDs = append(res.InsertedIDs, nil) } return &res, nil }