ITWorx-Pulse-Public/internal/metricquery/service.go

package metricquery

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"sync"
	"time"

	"github.com/itworx/pulse/internal/auth"
	"github.com/itworx/pulse/internal/metriccatalog"
	"github.com/itworx/pulse/internal/prometheus"
	"github.com/itworx/pulse/internal/promqlbinding"
	"github.com/itworx/pulse/internal/queryplan"
	"github.com/itworx/pulse/internal/redaction"
)

var ErrSourceUnavailable = errors.New("metric source unavailable")
var ErrInspectorUnauthorized = errors.New("metric inspector requires operate permission")

const (
	defaultCacheEntries = 128
	defaultCacheBytes   = 16 << 20
	defaultCacheTTL     = 15 * time.Second
)

// Freshness states carried by every metric response and live sample.
const (
	FreshnessFresh   = "fresh"
	FreshnessStale   = "stale"
	FreshnessUnknown = "unknown"
)

// FreshnessWindow is the freshness policy of a semantic metric definition.
func FreshnessWindow(definition metriccatalog.Definition) time.Duration {
	return time.Duration(definition.FreshnessSeconds) * time.Second
}

// ClassifyFreshness maps an observation timestamp onto the shared freshness
// vocabulary. Telemetry without a timestamp or without a policy is unknown and
// telemetry older than its policy window is stale; neither is ever fresh
// (ADR-0008).
func ClassifyFreshness(observedAt, now time.Time, maxAge time.Duration) string {
	if observedAt.IsZero() || maxAge <= 0 {
		return FreshnessUnknown
	}
	if now.Sub(observedAt) > maxAge {
		return FreshnessStale
	}
	return FreshnessFresh
}

type Source interface {
	Query(context.Context, string, *time.Time) (prometheus.QueryResult, error)
	QueryRange(context.Context, string, time.Time, time.Time, time.Duration) (prometheus.QueryResult, error)
}

type Provenance struct {
	Source         string `json:"source"`
	Metric         string `json:"metric"`
	CatalogVersion string `json:"catalogVersion"`
	CacheKey       string `json:"cacheKey"`
}

type InspectorLimits struct {
	MaxSeries int `json:"maxSeries"`
	MaxPoints int `json:"maxPoints"`
}

type Inspector struct {
	SemanticMetric string          `json:"semanticMetric"`
	GeneratedQuery string          `json:"generatedQuery"`
	Cost           queryplan.Cost  `json:"cost"`
	Limits         InspectorLimits `json:"limits"`
}

type Response struct {
	Status           string          `json:"status"`
	Data             json.RawMessage `json:"data"`
	Warnings         []string        `json:"warnings,omitempty"`
	Provenance       Provenance      `json:"provenance"`
	SourceObservedAt time.Time       `json:"sourceObservedAt"`
	ReceivedAt       time.Time       `json:"receivedAt"`
	Freshness        string          `json:"freshness"`
	CacheHit         bool            `json:"cacheHit"`
	Inspector        *Inspector      `json:"inspector,omitempty"`
}

type Service struct {
	planner  queryplan.Planner
	source   Source
	cache    *Cache
	mu       sync.Mutex
	inflight map[string]*call
	cacheTTL time.Duration
}

type call struct {
	done     chan struct{}
	response Response
	err      error
}

func NewService(planner queryplan.Planner, source Source, cache *Cache) *Service {
	if cache == nil {
		cache = NewCache(defaultCacheEntries, defaultCacheBytes)
	}
	return &Service{planner: planner, source: source, cache: cache, inflight: make(map[string]*call), cacheTTL: defaultCacheTTL}
}
func (s *Service) Inspect(ctx context.Context, request queryplan.Request) (Inspector, error) {
	if s == nil {
		return Inspector{}, ErrSourceUnavailable
	}
	principal, ok := principalFrom(ctx)
	if !ok || !auth.Allows(principal.Role, auth.PermissionOperate) {
		return Inspector{}, ErrInspectorUnauthorized
	}
	plan, err := s.planner.Plan(ctx, request)
	if err != nil {
		return Inspector{}, err
	}
	query, err := promqlbinding.CompilePlan(plan)
	if err != nil {
		return Inspector{}, err
	}
	inspector := inspectorFor(principal, plan, query)
	if inspector == nil {
		return Inspector{}, ErrInspectorUnauthorized
	}
	return *inspector, nil
}

func (s *Service) ExecuteRange(ctx context.Context, request queryplan.Request) (Response, error) {
	return s.execute(ctx, request, false, time.Time{})
}
func (s *Service) ExecuteInstant(ctx context.Context, request InstantRequest) (Response, error) {
	at := request.At
	if at.IsZero() {
		at = time.Now().UTC()
	}
	queryRequest := queryplan.Request{Metric: request.Metric, Scope: request.Scope, Aggregation: request.Aggregation, GroupBy: request.GroupBy, MaxSeries: request.MaxSeries, MaxPoints: 10, Range: queryplan.Range{From: at.Add(-time.Minute), To: at, StepSeconds: 60}}
	return s.execute(ctx, queryRequest, true, at.UTC())
}

type InstantRequest struct {
	Metric      string            `json:"metric"`
	Scope       map[string]string `json:"scope,omitempty"`
	At          time.Time         `json:"at,omitempty"`
	Aggregation string            `json:"aggregation,omitempty"`
	GroupBy     []string          `json:"groupBy,omitempty"`
	MaxSeries   int               `json:"maxSeries,omitempty"`
}

type sourceError struct{ err error }

func (e sourceError) Error() string { return ErrSourceUnavailable.Error() }
func (e sourceError) Unwrap() error { return e.err }
func (e sourceError) Is(target error) bool {
	return target == ErrSourceUnavailable || errors.Is(e.err, target)
}

func (s *Service) execute(ctx context.Context, request queryplan.Request, instant bool, at time.Time) (Response, error) {
	if s == nil || s.source == nil {
		return Response{}, ErrSourceUnavailable
	}
	plan, err := s.planner.Plan(ctx, request)
	if err != nil {
		return Response{}, err
	}
	query, err := promqlbinding.CompilePlan(plan)
	if err != nil {
		return Response{}, err
	}
	principal, _ := principalFrom(ctx)
	key := scopedKey(principal.Subject, plan.CacheKey)
	if cached, ok := s.cache.Get(key); ok {
		cached.CacheHit = true
		cached.Inspector = inspectorFor(principal, plan, query)
		return cached, nil
	}
	s.mu.Lock()
	if existing, ok := s.inflight[key]; ok {
		s.mu.Unlock()
		select {
		case <-existing.done:
			cachedResponse := existing.response
			cachedResponse.CacheHit = true
			cachedResponse.Inspector = inspectorFor(principal, plan, query)
			return cachedResponse, existing.err
		case <-ctx.Done():
			return Response{}, ctx.Err()
		}
	}
	current := &call{done: make(chan struct{})}
	s.inflight[key] = current
	s.mu.Unlock()
	response, err := s.fetch(ctx, plan, query, instant, at, key)
	if err == nil {
		response.Inspector = inspectorFor(principal, plan, query)
		cacheResponse := response
		cacheResponse.Inspector = nil
		s.cache.Put(key, cacheResponse, s.cacheTTL)
	}
	s.mu.Lock()
	current.response, current.err = response, err
	delete(s.inflight, key)
	close(current.done)
	s.mu.Unlock()
	return response, err
}

func (s *Service) fetch(ctx context.Context, plan queryplan.Plan, query string, instant bool, at time.Time, key string) (Response, error) {
	now := time.Now().UTC()
	var result prometheus.QueryResult
	var err error
	if instant {
		result, err = s.source.Query(ctx, query, &at)
	} else {
		result, err = s.source.QueryRange(ctx, query, plan.Request.Range.From, plan.Request.Range.To, time.Duration(plan.Request.Range.StepSeconds)*time.Second)
	}
	if err != nil {
		return Response{}, sourceError{err: err}
	}
	warnings := make([]string, 0, len(result.Warnings))
	for _, warning := range result.Warnings {
		warnings = append(warnings, redaction.String(warning))
	}
	return Response{Status: result.Status, Data: append(json.RawMessage(nil), result.Data...), Warnings: warnings, Provenance: Provenance{Source: "prometheus", Metric: plan.Metric.SemanticName, CatalogVersion: s.planner.CatalogVersion(), CacheKey: key}, SourceObservedAt: now, ReceivedAt: now, Freshness: ClassifyFreshness(now, now, FreshnessWindow(plan.Metric))}, nil
}

func inspectorFor(principal auth.Principal, plan queryplan.Plan, query string) *Inspector {
	if !auth.Allows(principal.Role, auth.PermissionOperate) {
		return nil
	}
	return &Inspector{SemanticMetric: plan.Metric.SemanticName, GeneratedQuery: redaction.String(query), Cost: plan.Cost, Limits: InspectorLimits{MaxSeries: plan.Request.MaxSeries, MaxPoints: plan.Request.MaxPoints}}
}

func principalFrom(ctx context.Context) (auth.Principal, bool) { return auth.PrincipalFromContext(ctx) }
func scopedKey(subject, planKey string) string {
	digest := sha256.Sum256([]byte(subject + "\x00" + planKey))
	return hex.EncodeToString(digest[:])
}

func cloneResponse(response Response) Response {
	copy := response
	copy.Data = append(json.RawMessage(nil), response.Data...)
	copy.Warnings = append([]string(nil), response.Warnings...)
	return copy
}

type cacheEntry struct {
	response  Response
	expiresAt time.Time
	lastUsed  time.Time
	size      int64
}
type Cache struct {
	mu         sync.Mutex
	entries    map[string]cacheEntry
	maxEntries int
	maxBytes   int64
	bytes      int64
}

func NewCache(maxEntries int, maxBytes int64) *Cache {
	if maxEntries <= 0 {
		maxEntries = defaultCacheEntries
	}
	if maxBytes <= 0 {
		maxBytes = defaultCacheBytes
	}
	return &Cache{entries: make(map[string]cacheEntry), maxEntries: maxEntries, maxBytes: maxBytes}
}
func (c *Cache) Get(key string) (Response, bool) {
	now := time.Now().UTC()
	c.mu.Lock()
	defer c.mu.Unlock()
	entry, ok := c.entries[key]
	if !ok {
		return Response{}, false
	}
	if !now.Before(entry.expiresAt) {
		c.bytes -= entry.size
		delete(c.entries, key)
		return Response{}, false
	}
	entry.lastUsed = now
	c.entries[key] = entry
	return cloneResponse(entry.response), true
}
func (c *Cache) Put(key string, response Response, ttl time.Duration) {
	if ttl <= 0 {
		ttl = defaultCacheTTL
	}
	size := int64(len(response.Data)) + int64(len(response.Warnings))*64 + 512
	if size > c.maxBytes {
		return
	}
	now := time.Now().UTC()
	c.mu.Lock()
	defer c.mu.Unlock()
	if old, ok := c.entries[key]; ok {
		c.bytes -= old.size
	}
	for len(c.entries) >= c.maxEntries || c.bytes+size > c.maxBytes {
		oldestKey := ""
		var oldest time.Time
		for candidateKey, candidate := range c.entries {
			if oldestKey == "" || candidate.lastUsed.Before(oldest) {
				oldestKey, oldest = candidateKey, candidate.lastUsed
			}
		}
		if oldestKey == "" {
			break
		}
		c.bytes -= c.entries[oldestKey].size
		delete(c.entries, oldestKey)
	}
	c.entries[key] = cacheEntry{response: cloneResponse(response), expiresAt: now.Add(ttl), lastUsed: now, size: size}
	c.bytes += size
}
func (c *Cache) Len() int     { c.mu.Lock(); defer c.mu.Unlock(); return len(c.entries) }
func (c *Cache) Bytes() int64 { c.mu.Lock(); defer c.mu.Unlock(); return c.bytes }