spb/internetpressure

1// Package batcher buffers measurements and traceroutes, and delivers them as gzip-compressed signed batches.2//3// Delivery rules (docs/PROBE-PROTOCOL.md § POST /batch): flush every batch_flush_seconds or when max_batch is4// reached; health block at most once per minute; network error / 5xx → spool to disk and back off 5 s → 5 min,5// drain oldest first when the server is back (re-signed with a fresh timestamp, bytes unchanged); 4xx → drop6// (except "401 skew": resync clock and retry once).7package batcher89import (10	"bytes"11	"compress/gzip"12	"context"13	"encoding/json"14	"errors"15	"log/slog"16	"sync"17	"sync/atomic"18	"time"1920	"internetpressure.io/probe-agent/internal/client"21	"internetpressure.io/probe-agent/internal/protocol"22	"internetpressure.io/probe-agent/internal/spool"23)2425// Defaults and limits.26const (27	DefaultFlushEvery = 10 * time.Second28	DefaultMaxBatch   = 50029	HealthEvery       = time.Minute30	BackoffMin        = 5 * time.Second31	BackoffMax        = 5 * time.Minute32	FinalFlushTimeout = 5 * time.Second33	drainPerCycle     = 2034)3536// Poster is the subset of client.Client used here (interface for tests).37type Poster interface {38	PostBatch(ctx context.Context, gz []byte, timeout time.Duration) (*protocol.BatchResponse, error)39}4041// Hooks are optional callbacks.42type Hooks struct {43	Health         func() *protocol.Health // builds the health block (called ≤ once/minute)44	OnResponse     func(*protocol.BatchResponse)45	OnFlushFailure func()46	OnDelivered    func(t time.Time)47}4849// Batcher is the queue + flusher.50type Batcher struct {51	probeID string52	version string53	poster  Poster54	spool   *spool.Spool55	hooks   Hooks56	log     *slog.Logger5758	flushEvery atomic.Int64 // ns59	maxBatch   atomic.Int326061	mu    sync.Mutex62	meas  []protocol.Measurement63	trs   []protocol.Traceroute64	kick  chan struct{}65	timer *time.Timer6667	lastHealth time.Time68	backoff    time.Duration69	nextTry    time.Time70	stopped    atomic.Bool71}7273// New creates a Batcher.74func New(probeID, version string, poster Poster, sp *spool.Spool, hooks Hooks, log *slog.Logger) *Batcher {75	if log == nil {76		log = slog.Default()77	}78	b := &Batcher{probeID: probeID, version: version, poster: poster, spool: sp, hooks: hooks, log: log,79		kick: make(chan struct{}, 1)}80	b.flushEvery.Store(int64(DefaultFlushEvery))81	b.maxBatch.Store(DefaultMaxBatch)82	return b83}8485// Configure updates flush interval / batch size (from the remote schedule). Zero keeps the default.86func (b *Batcher) Configure(flushSeconds, maxBatch int) {87	if flushSeconds > 0 {88		b.flushEvery.Store(int64(time.Duration(flushSeconds) * time.Second))89	}90	if maxBatch > 0 {91		b.maxBatch.Store(int32(maxBatch))92	}93}9495// Add queues measurements; triggers an immediate flush when max_batch is reached.96func (b *Batcher) Add(ms ...protocol.Measurement) {97	if len(ms) == 0 {98		return99	}100	b.mu.Lock()101	b.meas = append(b.meas, ms...)102	full := len(b.meas) >= int(b.maxBatch.Load())103	b.mu.Unlock()104	if full {105		b.Kick()106	}107}108109// AddTraceroute queues a traceroute.110func (b *Batcher) AddTraceroute(t protocol.Traceroute) {111	b.mu.Lock()112	b.trs = append(b.trs, t)113	b.mu.Unlock()114}115116// Kick requests a flush as soon as possible.117func (b *Batcher) Kick() {118	select {119	case b.kick <- struct{}{}:120	default:121	}122}123124// Buffered returns the number of queued measurements + traceroutes.125func (b *Batcher) Buffered() int {126	b.mu.Lock()127	defer b.mu.Unlock()128	return len(b.meas) + len(b.trs)129}130131// Run flushes on the timer and on kicks until ctx is cancelled. The caller then invokes Stop once the132// producers (scheduler) have finished, so the final flush contains every completed measurement.133func (b *Batcher) Run(ctx context.Context) {134	for {135		wait := time.Duration(b.flushEvery.Load())136		select {137		case <-ctx.Done():138			return139		case <-b.kick:140		case <-time.After(wait):141		}142		b.Flush(ctx)143	}144}145146// Stop performs the graceful final flush: one attempt with a 5 s timeout, everything else spooled.147func (b *Batcher) Stop() {148	if !b.stopped.CompareAndSwap(false, true) {149		return150	}151	ctx, cancel := context.WithTimeout(context.Background(), FinalFlushTimeout)152	defer cancel()153	for b.Buffered() > 0 {154		body, n := b.take()155		if n == 0 {156			break157		}158		gz, err := encode(body)159		if err != nil {160			b.log.Error("final flush: encode", "err", err)161			return162		}163		if _, err := b.poster.PostBatch(ctx, gz, FinalFlushTimeout); err != nil {164			b.toSpool(gz, "final flush failed: "+err.Error())165			// After one failure, spool the rest without trying the network again.166			for b.Buffered() > 0 {167				body, n := b.take()168				if n == 0 {169					break170				}171				if gz, err := encode(body); err == nil {172					b.toSpool(gz, "final flush: spooled remainder")173				}174			}175			return176		}177		b.log.Info("final flush delivered", "measurements", len(body.Measurements), "traceroutes", len(body.Traceroutes))178	}179}180181// Flush runs one delivery cycle: drain the spool (oldest first), then send the live queue.182func (b *Batcher) Flush(ctx context.Context) {183	now := time.Now()184	inBackoff := b.backoff > 0 && now.Before(b.nextTry)185186	if !inBackoff && b.spool != nil {187		b.drainSpool(ctx)188		inBackoff = b.backoff > 0 && time.Now().Before(b.nextTry)189	}190191	for {192		body, n := b.take()193		if n == 0 {194			return195		}196		gz, err := encode(body)197		if err != nil {198			b.log.Error("batch encode failed, dropping", "err", err, "n", n)199			return200		}201		if inBackoff {202			b.toSpool(gz, "server unavailable (backoff)")203		} else if !b.send(ctx, gz, body) {204			inBackoff = true205		}206		// Keep going until the queue is empty (in max_batch-sized chunks).207		b.mu.Lock()208		more := len(b.meas) > 0 || len(b.trs) > 0209		b.mu.Unlock()210		if !more {211			return212		}213	}214}215216// take removes up to max_batch measurements (and traceroutes) from the queue and builds the batch body,217// attaching the health block when due. Returns (body, number of items taken).218func (b *Batcher) take() (*protocol.Batch, int) {219	max := int(b.maxBatch.Load())220	b.mu.Lock()221	nm := len(b.meas)222	if nm > max {223		nm = max224	}225	nt := len(b.trs)226	if nt > max {227		nt = max228	}229	body := &protocol.Batch{ProbeID: b.probeID, AgentVersion: b.version,230		Measurements: make([]protocol.Measurement, nm)}231	copy(body.Measurements, b.meas[:nm])232	b.meas = append(b.meas[:0:0], b.meas[nm:]...) // fresh backing array so memory is released233	if nt > 0 {234		body.Traceroutes = make([]protocol.Traceroute, nt)235		copy(body.Traceroutes, b.trs[:nt])236		b.trs = append(b.trs[:0:0], b.trs[nt:]...)237	}238	b.mu.Unlock()239240	healthDue := b.hooks.Health != nil && time.Since(b.lastHealth) >= HealthEvery241	if nm+nt == 0 && !healthDue {242		return nil, 0243	}244	if healthDue {245		body.Health = b.hooks.Health()246		b.lastHealth = time.Now()247	}248	body.SentAt = protocol.FormatTime(time.Now())249	n := nm + nt250	if n == 0 {251		n = 1 // health-only batch252	}253	return body, n254}255256// encode marshals + gzips a batch.257func encode(body *protocol.Batch) ([]byte, error) {258	raw, err := json.Marshal(body)259	if err != nil {260		return nil, err261	}262	var buf bytes.Buffer263	zw, _ := gzip.NewWriterLevel(&buf, gzip.BestSpeed)264	if _, err := zw.Write(raw); err != nil {265		return nil, err266	}267	if err := zw.Close(); err != nil {268		return nil, err269	}270	return buf.Bytes(), nil271}272273// send posts one live batch. Returns false when the server is unavailable (batch spooled, backoff set).274func (b *Batcher) send(ctx context.Context, gz []byte, body *protocol.Batch) bool {275	resp, err := b.poster.PostBatch(ctx, gz, client.BatchTimeout)276	if err != nil {277		var he *client.HTTPError278		if errors.As(err, &he) && he.IsSkew() {279			// The client resynced its clock from the response; retry once with a fresh signature.280			b.log.Warn("batch rejected for clock skew, resynced and retrying once", "err", err)281			resp, err = b.poster.PostBatch(ctx, gz, client.BatchTimeout)282		}283	}284	if err != nil {285		if client.Retryable(err) {286			b.fail(err)287			b.toSpool(gz, err.Error())288			return false289		}290		b.log.Error("batch rejected, dropping", "err", err, "measurements", len(body.Measurements))291		return true292	}293	b.success(resp, len(body.Measurements), len(body.Traceroutes))294	return true295}296297// drainSpool re-sends spooled batches oldest first until empty, a failure, or drainPerCycle files.298func (b *Batcher) drainSpool(ctx context.Context) {299	for i := 0; i < drainPerCycle; i++ {300		name, gz, err := b.spool.Oldest()301		if err != nil {302			return303		}304		resp, err := b.poster.PostBatch(ctx, gz, client.BatchTimeout)305		if err != nil {306			var he *client.HTTPError307			if errors.As(err, &he) && he.IsSkew() {308				resp, err = b.poster.PostBatch(ctx, gz, client.BatchTimeout)309			}310		}311		if err != nil {312			if client.Retryable(err) {313				b.fail(err)314				return315			}316			b.log.Error("spooled batch rejected, dropping", "file", name, "err", err)317			b.spool.Remove(name)318			continue319		}320		b.spool.Remove(name)321		b.success(resp, -1, -1)322		b.log.Info("spooled batch delivered", "file", name, "accepted", resp.Accepted, "rejected", resp.Rejected)323	}324}325326func (b *Batcher) toSpool(gz []byte, why string) {327	if b.spool == nil {328		b.log.Error("no spool, batch lost", "why", why)329		return330	}331	name, err := b.spool.Write(gz)332	if err != nil {333		b.log.Error("spool write failed, batch lost", "err", err, "why", why)334		return335	}336	b.log.Warn("batch spooled", "file", name, "bytes", len(gz), "why", why)337}338339func (b *Batcher) fail(err error) {340	if b.backoff == 0 {341		b.backoff = BackoffMin342	} else {343		b.backoff *= 2344		if b.backoff > BackoffMax {345			b.backoff = BackoffMax346		}347	}348	b.nextTry = time.Now().Add(b.backoff)349	if b.hooks.OnFlushFailure != nil {350		b.hooks.OnFlushFailure()351	}352	b.log.Warn("batch delivery failed", "err", err, "retry_in", b.backoff)353}354355func (b *Batcher) success(resp *protocol.BatchResponse, nm, nt int) {356	if b.backoff > 0 {357		b.log.Info("server reachable again")358	}359	b.backoff = 0360	if b.hooks.OnDelivered != nil {361		b.hooks.OnDelivered(time.Now())362	}363	if resp != nil && b.hooks.OnResponse != nil {364		b.hooks.OnResponse(resp)365	}366	if nm >= 0 {367		b.log.Debug("batch delivered", "measurements", nm, "traceroutes", nt, "accepted", resp.Accepted, "rejected", resp.Rejected)368	}369}370371// InBackoff reports whether deliveries are currently suspended.372func (b *Batcher) InBackoff() bool { return b.backoff > 0 && time.Now().Before(b.nextTry) }373