128 lines
2.8 KiB
Go
128 lines
2.8 KiB
Go
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
|
|
}
|