pve-exporter/metrics/pve_ttl_metrics.go

package metrics

import (
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
)

// TTLMetric is an interface for metrics that support time-to-live cleanup.
// Any metric that implements Cleanup() can be registered in the TTLRegistry.
type TTLMetric interface {
	// Cleanup removes stale metric label sets that have exceeded their TTL.
	Cleanup()
}

// TTLGaugeVec wraps a Prometheus GaugeVec and tracks the last update time
// for each set of labels. When a set of labels is not updated within the TTL,
// it is automatically removed from the underlying GaugeVec.
type TTLGaugeVec struct {
	gaugeVec   *prometheus.GaugeVec // Underlying Prometheus GaugeVec.
	labelNames []string             // Label names of the underlying GaugeVec.
	ttl        time.Duration        // Duration after which an unused label set is considered stale.
	lastUpdate sync.Map             // Map storing last update time for each label set (key is a sorted labels string).
}

// NewTTLGaugeVec creates a new TTLGaugeVec using the provided GaugeOpts, label names, and TTL.
// The underlying GaugeVec is registered using promauto.
func NewTTLGaugeVec(opts prometheus.GaugeOpts, labelNames []string, ttl time.Duration) *TTLGaugeVec {
	return &TTLGaugeVec{
		gaugeVec:   promauto.NewGaugeVec(opts, labelNames),
		labelNames: labelNames,
		ttl:        ttl,
	}
}

// With returns the gauge for the given label set and records the current time
// as the last update for those labels.
func (t *TTLGaugeVec) With(labels prometheus.Labels) prometheus.Gauge {
	key := labelsKey(labels)
	t.lastUpdate.Store(key, time.Now())
	return t.gaugeVec.With(labels)
}

// Delete removes the metric associated with the given label set from both the underlying
// GaugeVec and the lastUpdate tracking map. It returns true if the deletion was successful.
func (t *TTLGaugeVec) Delete(labels prometheus.Labels) bool {
	key := labelsKey(labels)
	t.lastUpdate.Delete(key)
	return t.gaugeVec.Delete(labels)
}

// Cleanup iterates over all tracked label sets and deletes those that have not been updated
// within the TTL duration.
func (t *TTLGaugeVec) Cleanup() {
	now := time.Now()
	t.lastUpdate.Range(func(key, value interface{}) bool {
		if last, ok := value.(time.Time); ok {
			if now.Sub(last) > t.ttl {
				labels := parseLabels(key.(string))
				t.gaugeVec.Delete(labels)
				t.lastUpdate.Delete(key)
			}
		}
		return true
	})
}

// labelsKey creates a deterministic key from a Prometheus labels map.
// It sorts the keys and concatenates them in the format "key=value" separated by commas.
func labelsKey(labels prometheus.Labels) string {
	var keys []string
	for k := range labels {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	var parts []string
	for _, k := range keys {
		parts = append(parts, fmt.Sprintf("%s=%s", k, labels[k]))
	}
	return strings.Join(parts, ",")
}

// parseLabels converts a sorted key string back into a Prometheus labels map.
// The key is expected to be in the format produced by labelsKey.
func parseLabels(key string) prometheus.Labels {
	labels := prometheus.Labels{}
	parts := strings.Split(key, ",")
	for _, part := range parts {
		kv := strings.SplitN(part, "=", 2)
		if len(kv) == 2 {
			labels[kv[0]] = kv[1]
		}
	}
	return labels
}

// TTLRegistry manages multiple TTLMetric instances and periodically cleans them up.
type TTLRegistry struct {
	mu      sync.RWMutex
	metrics []TTLMetric
}

// NewTTLRegistry creates and returns a new TTLRegistry.
func NewTTLRegistry() *TTLRegistry {
	return &TTLRegistry{
		metrics: make([]TTLMetric, 0),
	}
}

// Register adds a TTLMetric to the registry for periodic cleanup.
func (r *TTLRegistry) Register(metric TTLMetric) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.metrics = append(r.metrics, metric)
}

// Cleanup calls the Cleanup method on each registered TTLMetric.
func (r *TTLRegistry) Cleanup() {
	r.mu.RLock()
	defer r.mu.RUnlock()
	for _, metric := range r.metrics {
		metric.Cleanup()
	}
}

// StartCleanupLoop starts a background goroutine that periodically cleans up stale metrics.
// The cleanup is performed at the specified interval.
func (r *TTLRegistry) StartCleanupLoop(interval time.Duration) {
	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		for range ticker.C {
			r.Cleanup()
		}
	}()
}

// TTLCounterVec exposes cumulative values as Prometheus counters while keeping
// the same "set an absolute value" usage as TTLGaugeVec. The PVE API reports
// counters as absolute totals, which prometheus.CounterVec cannot express
// (it only supports Inc/Add), so the metrics are emitted as constant metrics
// of type counter. Label sets not updated within the TTL are dropped.
type TTLCounterVec struct {
	desc       *prometheus.Desc            // Metric descriptor.
	labelNames []string                    // Label names in exposition order.
	ttl        time.Duration               // Duration after which an unused label set is considered stale.
	mu         sync.RWMutex                // Guards values.
	values     map[string]*ttlCounterEntry // Current value per label set.
}

// ttlCounterEntry holds the current value of a single label set.
type ttlCounterEntry struct {
	labelValues []string  // Label values in exposition order.
	value       float64   // Current counter value.
	lastUpdate  time.Time // Time of the last update.
}

// TTLCounter is a handle to a single label set of a TTLCounterVec.
type TTLCounter struct {
	vec   *TTLCounterVec
	entry *ttlCounterEntry
}

// Set stores the current absolute value of the counter.
func (c *TTLCounter) Set(value float64) {
	c.vec.mu.Lock()
	defer c.vec.mu.Unlock()
	c.entry.value = value
}

// NewTTLCounterVec creates a new TTLCounterVec and registers it using promauto.
func NewTTLCounterVec(opts prometheus.CounterOpts, labelNames []string, ttl time.Duration) *TTLCounterVec {
	c := &TTLCounterVec{
		desc: prometheus.NewDesc(
			prometheus.BuildFQName(opts.Namespace, opts.Subsystem, opts.Name),
			opts.Help,
			labelNames,
			opts.ConstLabels,
		),
		labelNames: labelNames,
		ttl:        ttl,
		values:     make(map[string]*ttlCounterEntry),
	}
	prometheus.MustRegister(c)
	return c
}

// With returns the counter handle for the given label set and records the current
// time as the last update for those labels.
func (t *TTLCounterVec) With(labels prometheus.Labels) *TTLCounter {
	key := labelsKey(labels)

	t.mu.Lock()
	defer t.mu.Unlock()

	entry, ok := t.values[key]
	if !ok {
		labelValues := make([]string, len(t.labelNames))
		for i, name := range t.labelNames {
			labelValues[i] = labels[name]
		}
		entry = &ttlCounterEntry{labelValues: labelValues}
		t.values[key] = entry
	}
	entry.lastUpdate = time.Now()

	return &TTLCounter{vec: t, entry: entry}
}

// Delete removes the metric associated with the given label set.
func (t *TTLCounterVec) Delete(labels prometheus.Labels) bool {
	key := labelsKey(labels)

	t.mu.Lock()
	defer t.mu.Unlock()

	if _, ok := t.values[key]; !ok {
		return false
	}
	delete(t.values, key)
	return true
}

// Cleanup deletes all label sets that have not been updated within the TTL duration.
func (t *TTLCounterVec) Cleanup() {
	now := time.Now()

	t.mu.Lock()
	defer t.mu.Unlock()

	for key, entry := range t.values {
		if now.Sub(entry.lastUpdate) > t.ttl {
			delete(t.values, key)
		}
	}
}

// Describe implements prometheus.Collector.
func (t *TTLCounterVec) Describe(ch chan<- *prometheus.Desc) {
	ch <- t.desc
}

// Collect implements prometheus.Collector.
func (t *TTLCounterVec) Collect(ch chan<- prometheus.Metric) {
	t.mu.RLock()
	defer t.mu.RUnlock()

	for _, entry := range t.values {
		ch <- prometheus.MustNewConstMetric(t.desc, prometheus.CounterValue, entry.value, entry.labelValues...)
	}
}