nmon/internal/vmsup/queue.go

package vmsup

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/v10networks/nmon/internal/clock"
	"github.com/v10networks/nmon/internal/metrics"
	"github.com/v10networks/nmon/internal/selfmon"
)

// ErrStoreDown is returned by queries, and by Queue.Flush, while the store is unhealthy.
var ErrStoreDown = errors.New("vmsup: store is down")

// Target is where a Queue or Querier finds a store. *Supervisor implements it.
type Target interface {
	URL(store string) string
	Healthy(store string) bool
}

// QueueOptions tunes a Queue; zero values take the defaults noted.
type QueueOptions struct {
	MaxSamples int           // 2_000_000: beyond this the oldest samples are dropped
	MaxAge     time.Duration // 60 s: samples older than this are dropped
	FlushEvery time.Duration // 1 s
	BatchSize  int           // 50_000 samples per POST
	Clock      clock.Clock   // clock.Real()
}

// Queue is the bounded in-memory write queue in front of one store. It
// implements metrics.Writer. Write never touches the network; Run (or Flush)
// posts batches to <store>/api/v1/import/prometheus.
type Queue struct {
	name   string // store name; also the "queue" and "path" label value
	target Target
	opts   QueueOptions
	clk    clock.Clock
	hc     *http.Client

	mu      sync.Mutex
	items   []*chunk
	samples int

	flushMu sync.Mutex // one flush at a time

	depth     selfmon.Gauge
	highWater selfmon.Gauge
	dropped   selfmon.Counter
}

// chunk is up to BatchSize samples, already encoded.
type chunk struct {
	body []byte
	n    int
	at   time.Time // enqueue time
}

// NewQueue returns a queue that writes to the named store ("hot" or "long") of
// target. If target is a *Supervisor, Supervisor.Stop flushes the queue.
func NewQueue(name string, target Target, opts QueueOptions) *Queue {
	if opts.MaxSamples <= 0 {
		opts.MaxSamples = 2_000_000
	}
	if opts.MaxAge <= 0 {
		opts.MaxAge = 60 * time.Second
	}
	if opts.FlushEvery <= 0 {
		opts.FlushEvery = time.Second
	}
	if opts.BatchSize <= 0 {
		opts.BatchSize = 50_000
	}
	if opts.Clock == nil {
		opts.Clock = clock.Real()
	}
	q := &Queue{
		name: name, target: target, opts: opts, clk: opts.Clock,
		hc:        &http.Client{Timeout: 30 * time.Second},
		depth:     selfmon.Default.Gauge("nmon_write_queue_depth", "queue", name),
		highWater: selfmon.Default.Gauge("nmon_write_queue_high_water", "queue", name),
		dropped:   selfmon.Default.Counter("nmon_dropped_samples_total", "path", name),
	}
	if sup, ok := target.(*Supervisor); ok {
		sup.register(q)
	}
	return q
}

// Write encodes s and enqueues it. It never blocks on the store and never
// fails because the store is down; samples that do not fit are dropped and
// counted in nmon_dropped_samples_total.
func (q *Queue) Write(_ context.Context, s []metrics.Sample) error {
	if len(s) == 0 {
		return nil
	}
	now := q.clk.Now()
	var chunks []*chunk
	for len(s) > 0 {
		n := min(len(s), q.opts.BatchSize)
		body := make([]byte, 0, n*96)
		for _, x := range s[:n] {
			body = metrics.AppendPromText(body, x)
		}
		chunks = append(chunks, &chunk{body: body, n: n, at: now})
		s = s[n:]
	}

	q.mu.Lock()
	dropped := q.expireLocked(now)
	for _, c := range chunks {
		if c.n > q.opts.MaxSamples {
			dropped += c.n // can never fit
			continue
		}
		for q.samples+c.n > q.opts.MaxSamples && len(q.items) > 0 {
			dropped += q.dropHeadLocked()
		}
		q.items = append(q.items, c)
		q.samples += c.n
	}
	q.gaugesLocked()
	q.mu.Unlock()
	q.count(dropped)
	return nil
}

// Depth is the number of samples waiting.
func (q *Queue) Depth() int { q.mu.Lock(); defer q.mu.Unlock(); return q.samples }

// Capacity is the most samples the queue holds (QueueOptions.MaxSamples).
func (q *Queue) Capacity() int { return q.opts.MaxSamples }

// Run flushes every FlushEvery until ctx ends. It does not flush on exit;
// Supervisor.Stop and callers use Flush for that.
func (q *Queue) Run(ctx context.Context) {
	t := q.clk.NewTicker(q.opts.FlushEvery)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C():
			_ = q.Flush(ctx)
		}
	}
}

// Flush drops expired samples, then posts as much as the store will take.
// It returns ErrStoreDown if the store is unhealthy, or the first POST error;
// unsent samples stay queued.
func (q *Queue) Flush(ctx context.Context) error {
	q.flushMu.Lock()
	defer q.flushMu.Unlock()
	for {
		q.mu.Lock()
		q.count(q.expireLocked(q.clk.Now()))
		q.gaugesLocked()
		batch := q.batchLocked()
		q.mu.Unlock()
		if len(batch) == 0 {
			return nil
		}
		if !q.target.Healthy(q.name) {
			return ErrStoreDown
		}
		retry, err := q.post(ctx, batch)
		if err != nil && retry {
			return err
		}
		// Sent, or rejected as bad data (retrying would never succeed): remove.
		var n int
		for _, c := range batch {
			n += c.n
		}
		q.mu.Lock()
		q.removeLocked(batch)
		q.gaugesLocked()
		q.mu.Unlock()
		if err != nil {
			q.count(n)
			return err
		}
	}
}

// batchLocked returns the leading chunks that together fit in BatchSize
// (always at least one chunk).
func (q *Queue) batchLocked() []*chunk {
	var n int
	for i, c := range q.items {
		if i > 0 && n+c.n > q.opts.BatchSize {
			return q.items[:i:i]
		}
		n += c.n
	}
	return q.items[:len(q.items):len(q.items)]
}

func (q *Queue) post(ctx context.Context, batch []*chunk) (retry bool, err error) {
	var body []byte
	if len(batch) == 1 {
		body = batch[0].body
	} else {
		var size int
		for _, c := range batch {
			size += len(c.body)
		}
		body = make([]byte, 0, size)
		for _, c := range batch {
			body = append(body, c.body...)
		}
	}
	u := q.target.URL(q.name)
	if u == "" {
		return true, ErrStoreDown
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, u+"/api/v1/import/prometheus", bytes.NewReader(body))
	if err != nil {
		return true, err
	}
	req.Header.Set("Content-Type", "text/plain")
	resp, err := q.hc.Do(req)
	if err != nil {
		return true, fmt.Errorf("vmsup: write to %s store: %w", q.name, err)
	}
	defer func() { _ = resp.Body.Close() }()
	msg, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
	switch {
	case resp.StatusCode/100 == 2:
		return false, nil
	case resp.StatusCode == http.StatusBadRequest:
		// VM rejects the whole batch as malformed; it will never be accepted.
		return false, fmt.Errorf("vmsup: %s store rejected batch: %s: %s", q.name, resp.Status, bytes.TrimSpace(msg))
	default:
		return true, fmt.Errorf("vmsup: write to %s store: %s: %s", q.name, resp.Status, bytes.TrimSpace(msg))
	}
}

// expireLocked drops chunks older than MaxAge from the head and returns the
// number of samples dropped.
func (q *Queue) expireLocked(now time.Time) int {
	var dropped int
	for len(q.items) > 0 && now.Sub(q.items[0].at) > q.opts.MaxAge {
		dropped += q.dropHeadLocked()
	}
	return dropped
}

func (q *Queue) dropHeadLocked() int {
	c := q.items[0]
	q.items[0] = nil
	q.items = q.items[1:]
	q.samples -= c.n
	return c.n
}

// removeLocked removes sent from the head. Chunks expired meanwhile are
// already gone, so only chunks still at the head (by identity) are removed.
func (q *Queue) removeLocked(sent []*chunk) {
	for _, c := range sent {
		if len(q.items) > 0 && q.items[0] == c {
			q.dropHeadLocked()
		}
	}
	if len(q.items) == 0 {
		q.items = nil // release the backing array
	}
}

func (q *Queue) gaugesLocked() {
	q.depth.Set(float64(q.samples))
	q.highWater.SetMax(float64(q.samples))
}

func (q *Queue) count(n int) {
	if n > 0 {
		q.dropped.Add(float64(n))
	}
}