logarchiver/internal/consumer/runner.go

// Package consumer binds the JetStream pull consumer and runs the archive loop.
//
// The core correctness property: a batch's messages are acknowledged ONLY after
// the batch has been sealed, uploaded to S3, and indexed in ClickHouse. If any
// of those fails the messages are Nak'd (with a backoff) and JetStream
// redelivers them, so nothing is lost on a sink outage. This sink-conditional
// acking is the main thing logarchiver does that a stock Vector NATS consumer
// cannot.
package consumer

import (
	"context"
	"errors"
	"log/slog"
	"time"

	"git.unkin.net/unkin/logarchiver/internal/batcher"
	"github.com/nats-io/nats.go/jetstream"
)

// Persister stores a ready batch durably. On success the caller acks.
type Persister interface {
	Store(ctx context.Context, batch *batcher.Batch) (StoreResult, error)
}

// StoreResult mirrors archiver.StoreResult (kept local to avoid an import cycle;
// the archiver's result is adapted at the call site).
type StoreResult struct {
	ObjectKey   string
	Events      int
	RawBytes    int64
	StoredBytes int64
}

// Metrics is the optional metrics surface for the loop.
type Metrics interface {
	MessagesFetched(n int)
	Acked(n int)
	BatchFlushed(trigger string)
	SetPending(n int)
}

// Runner drives the fetch → batch → persist → ack loop.
type Runner struct {
	cons       jetstream.Consumer
	batcher    *batcher.Batcher
	persist    Persister
	log        *slog.Logger
	metrics    Metrics
	fetchBatch int
	pollWait   time.Duration
	nakBackoff time.Duration
	drainTO    time.Duration
	nowFn      func() time.Time
}

// Options configures a Runner.
type Options struct {
	Consumer   jetstream.Consumer
	Batcher    *batcher.Batcher
	Persister  Persister
	Logger     *slog.Logger
	Metrics    Metrics
	FetchBatch int
	// PollWait bounds each Fetch and thus how often age-based flushes are checked.
	PollWait time.Duration
	// NakBackoff delays redelivery after a persist failure.
	NakBackoff time.Duration
	// DrainTimeout bounds the shutdown flush.
	DrainTimeout time.Duration
}

// NewRunner builds a Runner.
func NewRunner(o Options) *Runner {
	fetch := o.FetchBatch
	if fetch <= 0 {
		fetch = 512
	}
	poll := o.PollWait
	if poll <= 0 {
		poll = time.Second
	}
	nak := o.NakBackoff
	if nak <= 0 {
		nak = 10 * time.Second
	}
	drain := o.DrainTimeout
	if drain <= 0 {
		drain = 30 * time.Second
	}
	log := o.Logger
	if log == nil {
		log = slog.Default()
	}
	return &Runner{
		cons:       o.Consumer,
		batcher:    o.Batcher,
		persist:    o.Persister,
		log:        log,
		metrics:    o.Metrics,
		fetchBatch: fetch,
		pollWait:   poll,
		nakBackoff: nak,
		drainTO:    drain,
		nowFn:      time.Now,
	}
}

// Run loops until ctx is cancelled, then drains open batches before returning.
func (r *Runner) Run(ctx context.Context) error {
	r.log.Info("archive loop started",
		"fetch_batch", r.fetchBatch, "poll_wait", r.pollWait.String())
	for {
		if ctx.Err() != nil {
			return r.drain()
		}

		// Age-based flush before fetching more.
		r.flushBatches(ctx, r.batcher.DueByAge(r.nowFn()), "age")

		msgs, err := r.cons.Fetch(r.fetchBatch, jetstream.FetchMaxWait(r.pollWait))
		if err != nil {
			if errors.Is(err, context.Canceled) || ctx.Err() != nil {
				return r.drain()
			}
			r.log.Warn("fetch failed", "err", err)
			r.sleep(ctx, r.pollWait)
			continue
		}

		n := 0
		for msg := range msgs.Messages() {
			n++
			r.route(ctx, msg)
		}
		if ferr := msgs.Error(); ferr != nil && !errors.Is(ferr, context.Canceled) {
			r.log.Warn("fetch iteration error", "err", ferr)
		}
		if r.metrics != nil {
			r.metrics.MessagesFetched(n)
			r.metrics.SetPending(r.batcher.Pending())
		}
	}
}

// route decodes a message and adds it to the batcher, flushing if the batch
// becomes full.
func (r *Runner) route(ctx context.Context, msg jetstream.Msg) {
	meta := extractMeta(msg.Data())
	full := r.batcher.Add(batcher.Item{
		Subject:   msg.Subject(),
		Raw:       msg.Data(),
		Host:      meta.Host,
		Timestamp: meta.Timestamp,
		HasTS:     meta.Ok,
		Ack:       msg,
	})
	if full != nil {
		r.flush(ctx, full, "full")
	}
}

// flushBatches flushes a slice of batches with the given trigger label.
func (r *Runner) flushBatches(ctx context.Context, batches []*batcher.Batch, trigger string) {
	for _, b := range batches {
		r.flush(ctx, b, trigger)
	}
}

// flush persists a batch and, only on success, acks its messages. On failure it
// Naks with a backoff so JetStream redelivers.
func (r *Runner) flush(ctx context.Context, b *batcher.Batch, trigger string) {
	if len(b.Items) == 0 {
		return
	}
	res, err := r.persist.Store(ctx, b)
	if err != nil {
		r.log.Error("persist failed; batch will be redelivered",
			"subject", b.Subject, "events", len(b.Items), "trigger", trigger, "err", err)
		r.nakAll(b)
		return
	}
	acked := r.ackAll(b)
	if r.metrics != nil {
		r.metrics.Acked(acked)
		r.metrics.BatchFlushed(trigger)
	}
	r.log.Info("object archived",
		"subject", b.Subject, "object_key", res.ObjectKey,
		"events", res.Events, "raw_bytes", res.RawBytes, "stored_bytes", res.StoredBytes,
		"trigger", trigger)
}

func (r *Runner) ackAll(b *batcher.Batch) int {
	n := 0
	for _, it := range b.Items {
		if msg, ok := it.Ack.(jetstream.Msg); ok {
			if err := msg.Ack(); err != nil {
				r.log.Warn("ack failed", "subject", b.Subject, "err", err)
				continue
			}
			n++
		}
	}
	return n
}

func (r *Runner) nakAll(b *batcher.Batch) {
	for _, it := range b.Items {
		if msg, ok := it.Ack.(jetstream.Msg); ok {
			_ = msg.NakWithDelay(r.nakBackoff)
		}
	}
}

// drain flushes all open batches during shutdown with a bounded timeout.
func (r *Runner) drain() error {
	batches := r.batcher.Drain()
	if len(batches) == 0 {
		r.log.Info("archive loop stopped; nothing to drain")
		return nil
	}
	ctx, cancel := context.WithTimeout(context.Background(), r.drainTO)
	defer cancel()
	r.log.Info("draining open batches", "batches", len(batches))
	r.flushBatches(ctx, batches, "shutdown")
	return nil
}

func (r *Runner) sleep(ctx context.Context, d time.Duration) {
	t := time.NewTimer(d)
	defer t.Stop()
	select {
	case <-ctx.Done():
	case <-t.C:
	}
}