logmower-shipper/pkg/sender/sender.go

128 lines
2.8 KiB
Go
Raw Normal View History

2022-11-09 16:07:28 +00:00
package sender
import (
"context"
"fmt"
2022-11-09 16:07:28 +00:00
"log"
"time"
m "git.k-space.ee/k-space/logmower-shipper/pkg/mongo"
2022-11-12 22:19:19 +00:00
"github.com/jtagcat/util/batch"
"github.com/jtagcat/util/retry"
2022-11-09 16:07:28 +00:00
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
2022-11-11 15:52:30 +00:00
"k8s.io/apimachinery/pkg/util/wait"
2022-11-09 16:07:28 +00:00
)
2022-11-11 15:09:12 +00:00
var (
Simulate = false
MaxBatchItems = 1000
)
2022-11-09 16:07:28 +00:00
const (
2022-11-11 15:09:12 +00:00
MaxBatchTime = 5 * time.Second
2022-11-09 16:07:28 +00:00
)
2022-11-11 15:52:30 +00:00
// 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,
}
}
2022-11-09 16:07:28 +00:00
type Queue <-chan m.Record
func (queue Queue) Sender(db *mongo.Collection, metricsFilename string, cancelOnError func()) {
2022-11-09 16:07:28 +00:00
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:
}
}
}()
2022-11-12 22:19:19 +00:00
batch.Batch(MaxBatchItems, MaxBatchTime, queue, batched)
2022-11-09 16:07:28 +00:00
// 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)
2022-11-12 22:19:19 +00:00
err := retry.OnError(backoff(), func() (_ bool, _ error) {
2022-11-11 15:52:30 +00:00
result, err := insertManyWithSimulate(db, batch)
2022-11-11 15:52:30 +00:00
var succeedCount int
if result != nil {
succeedCount = len(result.InsertedIDs)
}
promShipperDbSent.WithLabelValues(metricsFilename).Add(float64(succeedCount))
2022-11-11 15:52:30 +00:00
if err == nil {
return
}
2022-11-09 16:07:28 +00:00
promShipperDbSendError.WithLabelValues(metricsFilename).Add(1)
2022-11-11 15:52:30 +00:00
batch = batch[succeedCount:]
2022-11-11 15:52:30 +00:00
if succeedCount == len(batch) {
return true, fmt.Errorf("all insertions in batch were successful, yet failure in database: %w", err)
}
2022-11-11 15:52:30 +00:00
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
}
2022-11-11 15:52:30 +00:00
2022-11-09 16:07:28 +00:00
}
}
func insertManyWithSimulate(db *mongo.Collection, batch []m.Record) (*mongo.InsertManyResult, error) {
2022-11-11 15:09:12 +00:00
if !Simulate {
var batchBson []interface{} // mongo does not like typing
for _, b := range batch {
batchBson = append(batchBson, b.ToBson())
}
tru := true
2022-11-11 15:09:12 +00:00
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
}