proxy-pool/internal/platform/metrics/capacity.go

package metrics

import (
	"errors"
	"fmt"
	"sync"

	"github.com/prometheus/client_golang/prometheus"

	controllerPool "proxy-pool/internal/controller/pool"
)

// CapacityCollector aggregates Controller leader-term inventory samples. An
// upstream ID is retained only in process memory for replacement and removal;
// it is never emitted as a Prometheus label.
type CapacityCollector struct {
	mu         sync.Mutex
	byUpstream map[string]capacityState

	reads           *prometheus.CounterVec
	managed         *prometheus.Desc
	available       *prometheus.Desc
	effective       *prometheus.Desc
	pending         *prometheus.Desc
	activeUpstreams *prometheus.Desc
}

type capacityState struct {
	sourceID  string
	managed   int
	available int64
	effective int64
	pending   int
}

var (
	_ controllerPool.CapacityObserver = (*CapacityCollector)(nil)
	_ prometheus.Collector            = (*CapacityCollector)(nil)
)

func NewCapacityCollector(registerer prometheus.Registerer) (*CapacityCollector, error) {
	if registerer == nil {
		return nil, ErrInvalidDependencies
	}
	reads, err := registerCounterVec(registerer, prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "proxy_pool", Subsystem: "controller", Name: "capacity_inventory_reads_total",
		Help: "Number of Provider inventory reads by fixed result.",
	}, []string{"result"}))
	if err != nil {
		return nil, err
	}
	candidate := &CapacityCollector{
		byUpstream: map[string]capacityState{}, reads: reads,
		managed: prometheus.NewDesc(
			"proxy_pool_controller_capacity_managed_proxies",
			"Managed non-extracted proxies from the latest successful samples of active Controller Provider terms.", nil, nil,
		),
		available: prometheus.NewDesc(
			"proxy_pool_controller_capacity_available_slots",
			"Available proxy concurrency slots from the latest successful samples of active Controller Provider terms.", nil, nil,
		),
		effective: prometheus.NewDesc(
			"proxy_pool_controller_capacity_effective_slots",
			"Available slots plus expected pending fetch capacity from the latest successful active Controller samples.", nil, nil,
		),
		pending: prometheus.NewDesc(
			"proxy_pool_controller_capacity_pending_expected_proxies",
			"Expected proxy candidates reserved by in-flight fetches in the latest successful active Controller samples.", nil, nil,
		),
		activeUpstreams: prometheus.NewDesc(
			"proxy_pool_controller_capacity_active_upstreams",
			"Number of Upstreams with a successful inventory sample in this Controller process.", nil, nil,
		),
	}
	return registerCapacityCollector(registerer, candidate)
}

func (collector *CapacityCollector) ObserveCapacity(observation controllerPool.CapacityObservation) {
	if collector == nil || collector.reads == nil || observation.UpstreamID == "" || observation.SourceID == "" {
		return
	}
	switch observation.Result {
	case controllerPool.CapacityReadError:
		collector.reads.WithLabelValues(string(observation.Result)).Inc()
		return
	case controllerPool.CapacityReadSuccess:
		if observation.Managed < 0 || observation.AvailableSlots < 0 || observation.EffectiveSlots < 0 || observation.PendingExpected < 0 {
			return
		}
		collector.reads.WithLabelValues(string(observation.Result)).Inc()
	default:
		return
	}
	collector.mu.Lock()
	collector.byUpstream[observation.UpstreamID] = capacityState{
		sourceID: observation.SourceID,
		managed:  observation.Managed, available: observation.AvailableSlots,
		effective: observation.EffectiveSlots, pending: observation.PendingExpected,
	}
	collector.mu.Unlock()
}

func (collector *CapacityCollector) RemoveCapacityUpstream(upstreamID, sourceID string) {
	if collector == nil || upstreamID == "" || sourceID == "" {
		return
	}
	collector.mu.Lock()
	if state, exists := collector.byUpstream[upstreamID]; exists && state.sourceID == sourceID {
		delete(collector.byUpstream, upstreamID)
	}
	collector.mu.Unlock()
}

func (collector *CapacityCollector) Describe(descriptions chan<- *prometheus.Desc) {
	if collector == nil {
		return
	}
	descriptions <- collector.managed
	descriptions <- collector.available
	descriptions <- collector.effective
	descriptions <- collector.pending
	descriptions <- collector.activeUpstreams
}

func (collector *CapacityCollector) Collect(metrics chan<- prometheus.Metric) {
	if collector == nil {
		return
	}
	collector.mu.Lock()
	var managed, available, effective, pending int64
	for _, state := range collector.byUpstream {
		managed += int64(state.managed)
		available += state.available
		effective += state.effective
		pending += int64(state.pending)
	}
	activeUpstreams := len(collector.byUpstream)
	collector.mu.Unlock()
	metrics <- prometheus.MustNewConstMetric(collector.managed, prometheus.GaugeValue, float64(managed))
	metrics <- prometheus.MustNewConstMetric(collector.available, prometheus.GaugeValue, float64(available))
	metrics <- prometheus.MustNewConstMetric(collector.effective, prometheus.GaugeValue, float64(effective))
	metrics <- prometheus.MustNewConstMetric(collector.pending, prometheus.GaugeValue, float64(pending))
	metrics <- prometheus.MustNewConstMetric(collector.activeUpstreams, prometheus.GaugeValue, float64(activeUpstreams))
}

func registerCapacityCollector(registerer prometheus.Registerer, candidate *CapacityCollector) (*CapacityCollector, error) {
	if err := registerer.Register(candidate); err == nil {
		return candidate, nil
	} else {
		var registered prometheus.AlreadyRegisteredError
		if !errors.As(err, &registered) {
			return nil, fmt.Errorf("register capacity collector: %w", err)
		}
		existing, ok := registered.ExistingCollector.(*CapacityCollector)
		if !ok {
			return nil, fmt.Errorf("register capacity collector: existing collector has unexpected type")
		}
		return existing, nil
	}
}