nmon/internal/vmsup/query.go

package vmsup

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

// maxQueryResponse bounds what a Querier reads into memory.
const maxQueryResponse = 64 << 20

// Querier runs PromQL/MetricsQL queries against a store. Calls fail fast with
// ErrStoreDown when the store is unhealthy instead of hanging on a dead socket.
type Querier struct {
	target Target
	hc     *http.Client
}

// NewQuerier returns a Querier for target (usually the *Supervisor).
func NewQuerier(target Target) *Querier {
	return &Querier{target: target, hc: &http.Client{Timeout: 60 * time.Second}}
}

// Query runs an instant query at ts (zero means now) and returns the raw JSON response.
func (c *Querier) Query(ctx context.Context, store, q string, ts time.Time) ([]byte, error) {
	form := url.Values{"query": {q}}
	if !ts.IsZero() {
		form.Set("time", unixSeconds(ts))
	}
	return c.do(ctx, store, "/api/v1/query", form)
}

// QueryRange runs a range query and returns the raw JSON response.
func (c *Querier) QueryRange(ctx context.Context, store, q string, from, to time.Time, step time.Duration) ([]byte, error) {
	if step <= 0 {
		return nil, fmt.Errorf("vmsup: query step must be positive, got %s", step)
	}
	form := url.Values{
		"query": {q},
		"start": {unixSeconds(from)},
		"end":   {unixSeconds(to)},
		"step":  {strconv.FormatFloat(step.Seconds(), 'f', -1, 64)},
	}
	return c.do(ctx, store, "/api/v1/query_range", form)
}

func (c *Querier) do(ctx context.Context, store, path string, form url.Values) ([]byte, error) {
	if !c.target.Healthy(store) {
		return nil, fmt.Errorf("%w: %s", ErrStoreDown, store)
	}
	base := c.target.URL(store)
	if base == "" {
		return nil, fmt.Errorf("%w: %s", ErrStoreDown, store)
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, base+path, strings.NewReader(form.Encode()))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	resp, err := c.hc.Do(req)
	if err != nil {
		return nil, fmt.Errorf("vmsup: query %s store: %w", store, err)
	}
	defer func() { _ = resp.Body.Close() }()
	body, err := io.ReadAll(io.LimitReader(resp.Body, maxQueryResponse+1))
	if err != nil {
		return nil, fmt.Errorf("vmsup: read %s store response: %w", store, err)
	}
	if len(body) > maxQueryResponse {
		return nil, fmt.Errorf("vmsup: %s store response exceeds %d bytes", store, maxQueryResponse)
	}
	if resp.StatusCode != http.StatusOK {
		msg := body
		if len(msg) > 512 {
			msg = msg[:512]
		}
		return nil, fmt.Errorf("vmsup: %s store: %s: %s", store, resp.Status, strings.TrimSpace(string(msg)))
	}
	return body, nil
}

func unixSeconds(t time.Time) string {
	return strconv.FormatFloat(float64(t.UnixMilli())/1000, 'f', 3, 64)
}