logarchiver/internal/batcher/batcher.go

// Package batcher groups incoming log events into per-subject batches and
// decides when a batch is ready to become one archived object. It is
// deliberately not concurrent: the consumer run loop owns a Batcher and drives
// it from a single goroutine (Add on receive, DueByAge on a ticker, Drain on
// shutdown), which keeps the ack-after-persist accounting simple and race-free.
package batcher

import (
	"bytes"
	"sort"
	"time"
)

// Item is one log event routed into a batch. Ack is an opaque token (a
// jetstream.Msg in production) that the caller acknowledges only after the batch
// has been durably persisted.
type Item struct {
	Subject   string
	Raw       []byte
	Host      string
	Timestamp time.Time
	HasTS     bool
	Ack       any
}

// Limits bound a single batch/object.
type Limits struct {
	MaxBytes  int64
	MaxEvents int
	MaxAge    time.Duration
}

// Batch is a ready (or in-progress) group of events for one subject.
type Batch struct {
	Subject  string
	Items    []Item
	RawBytes int64
	OpenedAt time.Time
}

// Batcher accumulates open batches keyed by subject.
type Batcher struct {
	limits Limits
	open   map[string]*Batch
	nowFn  func() time.Time
}

// New returns a Batcher enforcing limits.
func New(limits Limits) *Batcher {
	return &Batcher{limits: limits, open: map[string]*Batch{}, nowFn: time.Now}
}

// Add appends it to its subject's open batch. If that batch is now full (by
// bytes or event count), it is removed from the open set and returned so the
// caller can flush it; otherwise Add returns nil.
func (b *Batcher) Add(it Item) *Batch {
	batch := b.open[it.Subject]
	if batch == nil {
		batch = &Batch{Subject: it.Subject, OpenedAt: b.nowFn()}
		b.open[it.Subject] = batch
	}
	batch.Items = append(batch.Items, it)
	batch.RawBytes += int64(len(it.Raw))

	if b.full(batch) {
		delete(b.open, it.Subject)
		return batch
	}
	return nil
}

func (b *Batcher) full(batch *Batch) bool {
	if b.limits.MaxBytes > 0 && batch.RawBytes >= b.limits.MaxBytes {
		return true
	}
	if b.limits.MaxEvents > 0 && len(batch.Items) >= b.limits.MaxEvents {
		return true
	}
	return false
}

// DueByAge removes and returns every open batch older than MaxAge as of now.
func (b *Batcher) DueByAge(now time.Time) []*Batch {
	if b.limits.MaxAge <= 0 {
		return nil
	}
	var due []*Batch
	for subj, batch := range b.open {
		if now.Sub(batch.OpenedAt) >= b.limits.MaxAge {
			due = append(due, batch)
			delete(b.open, subj)
		}
	}
	sortBatches(due)
	return due
}

// Drain removes and returns all open batches (used on graceful shutdown so
// in-flight events are persisted and acked before exit).
func (b *Batcher) Drain() []*Batch {
	var all []*Batch
	for subj, batch := range b.open {
		all = append(all, batch)
		delete(b.open, subj)
	}
	sortBatches(all)
	return all
}

// Pending reports how many events sit in open batches.
func (b *Batcher) Pending() int {
	n := 0
	for _, batch := range b.open {
		n += len(batch.Items)
	}
	return n
}

func sortBatches(bs []*Batch) {
	sort.Slice(bs, func(i, j int) bool { return bs[i].Subject < bs[j].Subject })
}

// NDJSON renders the batch as newline-delimited JSON (one raw event per line),
// matching the raw archive format the logging stack expects.
func (b *Batch) NDJSON() []byte {
	var buf bytes.Buffer
	buf.Grow(int(b.RawBytes) + len(b.Items))
	for _, it := range b.Items {
		buf.Write(bytes.TrimRight(it.Raw, "\n"))
		buf.WriteByte('\n')
	}
	return buf.Bytes()
}

// Summary is the index-relevant projection of a batch.
type Summary struct {
	Hosts      []string
	MinTS      time.Time
	MaxTS      time.Time
	EventCount int
	HasTS      bool
}

// Summarize computes hosts (unique, sorted) and the timestamp range. fallback is
// used for events whose payload lacked a parseable timestamp (ingest time).
func (b *Batch) Summarize(fallback time.Time) Summary {
	s := Summary{EventCount: len(b.Items)}
	hostSet := map[string]struct{}{}
	for _, it := range b.Items {
		if it.Host != "" {
			hostSet[it.Host] = struct{}{}
		}
		ts := it.Timestamp
		if !it.HasTS {
			ts = fallback
		} else {
			s.HasTS = true
		}
		if s.MinTS.IsZero() || ts.Before(s.MinTS) {
			s.MinTS = ts
		}
		if s.MaxTS.IsZero() || ts.After(s.MaxTS) {
			s.MaxTS = ts
		}
	}
	if s.MinTS.IsZero() {
		s.MinTS = fallback
	}
	if s.MaxTS.IsZero() {
		s.MaxTS = fallback
	}
	for h := range hostSet {
		s.Hosts = append(s.Hosts, h)
	}
	sort.Strings(s.Hosts)
	return s
}