ITWorx-Pulse-Public/internal/prometheus/client.go

package prometheus

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/itworx/pulse/internal/datasource"
)

type Limits struct {
	Timeout        time.Duration
	MaxConcurrency int
	MaxQueryLength int
	MaxResponse    int64
}

func (l Limits) withDefaults() Limits {
	if l.Timeout == 0 {
		l.Timeout = 10 * time.Second
	}
	if l.MaxConcurrency == 0 {
		l.MaxConcurrency = 4
	}
	if l.MaxQueryLength == 0 {
		l.MaxQueryLength = 4096
	}
	if l.MaxResponse == 0 {
		l.MaxResponse = 4 << 20
	}
	return l
}

// Metrics is the bounded adapter counter set. TotalDuration is the sum of the
// upstream round-trip durations, so the (Requests, TotalDuration) pair exposes
// mean query latency the same way a Prometheus count/sum histogram does, and
// MaxDuration keeps the worst observed round trip visible to operators.
type Metrics struct {
	Requests, Errors, Timeouts uint64
	TotalDuration, MaxDuration time.Duration
}

// MetricsSink is the narrow view of the process metrics registry the adapter
// needs. *observability.Registry satisfies it.
type MetricsSink interface {
	SetGauge(name string, value float64)
}

type Client struct {
	baseURL *url.URL
	http    *http.Client
	limits  Limits
	sem     chan struct{}
	mu      sync.Mutex
	metrics Metrics
}

func New(baseURL string, httpClient *http.Client, limits Limits) (*Client, error) {
	u, err := url.Parse(strings.TrimSpace(baseURL))
	if err != nil || u.Scheme != "http" && u.Scheme != "https" || u.Host == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" {
		return nil, errors.New("prometheus base URL must be an absolute HTTP(S) URL without credentials or query")
	}
	limits = limits.withDefaults()
	if limits.Timeout <= 0 || limits.Timeout > time.Minute || limits.MaxConcurrency < 1 || limits.MaxConcurrency > 32 || limits.MaxQueryLength < 1 || limits.MaxResponse < 1 {
		return nil, errors.New("invalid Prometheus adapter limits")
	}
	if httpClient == nil {
		httpClient = &http.Client{}
	}
	return &Client{baseURL: u, http: httpClient, limits: limits, sem: make(chan struct{}, limits.MaxConcurrency)}, nil
}

type QueryResult struct {
	Status    string          `json:"status"`
	Data      json.RawMessage `json:"data"`
	Warnings  []string        `json:"warnings"`
	ErrorType string          `json:"errorType"`
	Error     string          `json:"error"`
}

func (c *Client) Query(ctx context.Context, query string, at *time.Time) (QueryResult, error) {
	if len(query) == 0 || len(query) > c.limits.MaxQueryLength {
		return QueryResult{}, errors.New("PromQL query exceeds bounds")
	}
	values := url.Values{"query": []string{query}}
	if at != nil {
		values.Set("time", at.UTC().Format(time.RFC3339Nano))
	}
	return c.request(ctx, "/api/v1/query", values)
}

func (c *Client) QueryRange(ctx context.Context, query string, start, end time.Time, step time.Duration) (QueryResult, error) {
	if len(query) == 0 || len(query) > c.limits.MaxQueryLength || !start.Before(end) || step <= 0 || step > 24*time.Hour {
		return QueryResult{}, errors.New("invalid bounded Prometheus range query")
	}
	values := url.Values{"query": []string{query}, "start": []string{start.UTC().Format(time.RFC3339Nano)}, "end": []string{end.UTC().Format(time.RFC3339Nano)}, "step": []string{step.String()}}
	return c.request(ctx, "/api/v1/query_range", values)
}

func (c *Client) request(ctx context.Context, path string, values url.Values) (QueryResult, error) {
	select {
	case c.sem <- struct{}{}:
	case <-ctx.Done():
		return QueryResult{}, ctx.Err()
	}
	defer func() { <-c.sem }()
	requestContext, cancel := context.WithTimeout(ctx, c.limits.Timeout)
	defer cancel()
	u := *c.baseURL
	u.Path = strings.TrimRight(u.Path, "/") + path
	u.RawQuery = values.Encode()
	req, err := http.NewRequestWithContext(requestContext, http.MethodGet, u.String(), nil)
	if err != nil {
		return QueryResult{}, err
	}
	started := time.Now()
	response, err := c.http.Do(req)
	elapsed := time.Since(started)
	c.mu.Lock()
	c.metrics.Requests++
	c.metrics.TotalDuration += elapsed
	if elapsed > c.metrics.MaxDuration {
		c.metrics.MaxDuration = elapsed
	}
	if err != nil {
		c.metrics.Errors++
		if errors.Is(err, context.DeadlineExceeded) {
			c.metrics.Timeouts++
		}
	}
	c.mu.Unlock()
	if err != nil {
		return QueryResult{}, fmt.Errorf("prometheus request: %w", err)
	}
	defer response.Body.Close()
	if response.ContentLength > c.limits.MaxResponse {
		return QueryResult{}, errors.New("prometheus response exceeds bounds")
	}
	body, err := io.ReadAll(io.LimitReader(response.Body, c.limits.MaxResponse+1))
	if err != nil || int64(len(body)) > c.limits.MaxResponse {
		return QueryResult{}, errors.New("prometheus response exceeds bounds")
	}
	var result QueryResult
	if err := json.Unmarshal(body, &result); err != nil {
		return QueryResult{}, errors.New("invalid Prometheus response")
	}
	if response.StatusCode < 200 || response.StatusCode >= 300 || result.Status != "success" {
		return result, fmt.Errorf("prometheus upstream error: %s", strings.TrimSpace(result.ErrorType+" "+result.Error))
	}
	return result, nil
}

func (c *Client) Metrics() Metrics { c.mu.Lock(); defer c.mu.Unlock(); return c.metrics }

// PublishMetrics copies the adapter counters and query latency onto the process
// metrics registry that backs GET /api/v1/system/metrics. Names are fixed, so
// the registry stays bounded no matter how many queries are served.
func (c *Client) PublishMetrics(sink MetricsSink) {
	if c == nil || sink == nil {
		return
	}
	metrics := c.Metrics()
	sink.SetGauge("pulse_prometheus_requests_total", float64(metrics.Requests))
	sink.SetGauge("pulse_prometheus_errors_total", float64(metrics.Errors))
	sink.SetGauge("pulse_prometheus_timeouts_total", float64(metrics.Timeouts))
	sink.SetGauge("pulse_prometheus_request_duration_seconds_sum", metrics.TotalDuration.Seconds())
	sink.SetGauge("pulse_prometheus_request_duration_seconds_max", metrics.MaxDuration.Seconds())
}

func (c *Client) Health(ctx context.Context) datasource.SourceHealth {
	now := time.Now().UTC()
	_, err := c.Query(ctx, "up", nil)
	state := datasource.HealthHealthy
	reason := ""
	lastSuccess := now
	if err != nil {
		state = datasource.HealthUnknown
		reason = "PROMETHEUS_UNAVAILABLE"
		lastSuccess = time.Time{}
	}
	return datasource.SourceHealth{State: state, ObservedAt: now, ReceivedAt: time.Now().UTC(), LastSuccess: lastSuccess, Policy: datasource.FreshnessPolicy{MaxAge: 30 * time.Second}, ReasonCode: reason}
}