logmower-shipper/pkg/sender/sender.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
}