stern-event-agent/shipper/shipper.go

package main

// The pure half of the shipper: line parsing, the Loki push payload, the offset state.
// main.go is the tail-and-push loop around it.

import (
	"encoding/json"
	"os"
	"regexp"
	"strconv"
	"time"
)

const stampLayout = "2006-01-02T15:04:05.000Z"

var stampRe = regexp.MustCompile(`\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{3}Z`)
var gameRe = regexp.MustCompile(`up: game (\S+)`)

// extractTimestamp finds the line's own RFC 3339 (millisecond, UTC) stamp wherever it sits:
// the runtime prefixes an uptime count and a [who], gamelog puts its stamp after its tag,
// the switch log starts with it. A line without a stamp (the boot banner) reports false and
// the caller stamps it with the shipping time.
func extractTimestamp(line string) (time.Time, bool) {
	s := stampRe.FindString(line)
	if s == "" {
		return time.Time{}, false
	}
	ts, err := time.Parse(stampLayout, s)
	if err != nil {
		return time.Time{}, false
	}
	return ts, true
}

// extractGame names the game from the boot line ("up: game king_kong_le 0.97, ..."), "" elsewhere.
func extractGame(line string) string {
	m := gameRe.FindStringSubmatch(line)
	if m == nil {
		return ""
	}
	return m[1]
}

// assignTimestamps gives every line a timestamp: its own stamp where it has one; an
// unstamped line (the boot banner) takes the NEXT stamped line's time — a stream must
// not start at the shipping time and then jump back to the game's boot — and a trailing
// unstamped line takes the previous stamped time. fallback covers a batch with no
// stamps at all.
func assignTimestamps(lines []string, fallback time.Time) []time.Time {
	stamps := make([]time.Time, len(lines))
	next := time.Time{}
	for i := len(lines) - 1; i >= 0; i-- {
		if ts, ok := extractTimestamp(lines[i]); ok {
			next = ts
		}
		stamps[i] = next
	}
	prev := time.Time{}
	for i := range lines {
		if ts, ok := extractTimestamp(lines[i]); ok {
			prev = ts
		}
		if stamps[i].IsZero() {
			stamps[i] = prev
		}
	}
	for i := range stamps {
		if stamps[i].IsZero() {
			stamps[i] = fallback
		}
	}
	return stamps
}

// Entry is one shipped line.
type Entry struct {
	Source string
	Game   string
	TS     time.Time
	Line   string
}

// lokiPush is the /loki/api/v1/push body: labelled streams of (nanoseconds, line) pairs.
type lokiPush struct {
	Streams []lokiStream `json:"streams"`
}

type lokiStream struct {
	Stream map[string]string `json:"stream"`
	Values [][]string        `json:"values"`
}

// buildPayloads groups entries by their label set and marshals Loki push bodies, each JSON
// under maxBytes (Loki refuses a push past its limit, so a big batch is split).
func buildPayloads(entries []Entry, machine string, maxBytes int) [][]byte {
	type group struct {
		labels map[string]string
		values [][]string
	}
	groups := map[string]*group{}
	var order []string
	for _, e := range entries {
		ts, ok := extractTimestamp(e.Line)
		if !ok {
			ts = e.TS
		}
		labels := map[string]string{"job": "pinball", "machine": machine, "source": e.Source}
		if e.Game != "" {
			labels["game"] = e.Game
		}
		key := e.Source + "\x00" + e.Game
		g, ok := groups[key]
		if !ok {
			groups[key] = &group{labels: labels}
			g = groups[key]
			order = append(order, key)
		}
		g.values = append(g.values, []string{strconv.FormatInt(ts.UnixNano(), 10), e.Line})
	}

	var payloads [][]byte
	var cur *lokiPush
	curBytes := 0
	flush := func() {
		if cur != nil && len(cur.Streams) > 0 {
			if b, err := json.Marshal(cur); err == nil {
				payloads = append(payloads, b)
			}
		}
		cur = nil
		curBytes = 0
	}
	add := func(s lokiStream) int {
		p := lokiPush{Streams: []lokiStream{s}}
		b, err := json.Marshal(p)
		if err != nil {
			return 0
		}
		if cur == nil {
			cur = &lokiPush{}
		}
		cur.Streams = append(cur.Streams, s)
		curBytes += len(b)
		return len(b)
	}
	for _, key := range order {
		g := groups[key]
		for len(g.values) > 0 {
			// a stream that fits whole goes into the current payload; one that does not is
			// halved until a piece fits, then the payload is flushed and the rest continues
			take := g.values
			var piece []lokiStream
			for {
				s := lokiStream{Stream: g.labels, Values: take}
				p := lokiPush{Streams: []lokiStream{s}}
				b, err := json.Marshal(p)
				if err != nil {
					break
				}
				if len(b)+curBytes <= maxBytes {
					piece = []lokiStream{s}
					break
				}
				if len(take) == 1 {
					// a single line over the cap is shipped alone anyway (it cannot split)
					piece = []lokiStream{s}
					break
				}
				take = take[:len(take)/2]
			}
			if piece == nil {
				flush()
				continue
			}
			add(piece[0])
			g.values = g.values[len(piece[0].Values):]
			if len(g.values) > 0 || curBytes > maxBytes/2 {
				flush()
			}
		}
	}
	flush()
	return payloads
}

// needsFlush: entries older than the out-of-order window (Loki serves head entries only
// within max_chunk_age/2 ≈ 1h) sit invisible in the ingester until a flush puts them in the
// store. A backfill (first ship of history) needs one; a live tail never does.
func needsFlush(entries []Entry, now time.Time) bool {
	for _, e := range entries {
		if now.Sub(e.TS) > time.Hour {
			return true
		}
	}
	return false
}

// ---- the offset state -------------------------------------------------------------------

type state map[string]int64

// loadState reads the saved per-file offsets; a missing file is a fresh start, not an error.
func loadState(path string) (state, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			return state{}, nil
		}
		return state{}, err
	}
	var st state
	if err := json.Unmarshal(b, &st); err != nil {
		return state{}, err
	}
	return st, nil
}

// saveState writes the offsets atomically (a tmp file and a rename), so a crash mid-write
// keeps the last fully saved state.
func saveState(path string, st state) error {
	b, err := json.Marshal(st)
	if err != nil {
		return err
	}
	tmp := path + ".tmp"
	if err := os.WriteFile(tmp, b, 0644); err != nil {
		return err
	}
	return os.Rename(tmp, path)
}