// SPDX-License-Identifier: Apache-2.0
//
// Copyright The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package metriccache reuses const metrics across scrapes for objects whose
// values rarely change.
//
// Building a const metric allocates its label pairs and value, which dominates
// the allocations of collectors that publish many state-set series per object.
// Const metrics are immutable and a registry only reads them during Gather, so
// sending the same metric again in a later scrape is safe.
//
// A collector lists its caches in ScrapeCaches (see internal/scrapecache).
// The collection resets them around Build and Close and begins and ends a
// scrape around each Collect call, which gets it from [Cache.Scrape].
package metriccache

import (
	"sync"
	"sync/atomic"

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

// Cache holds the metrics built for each key during the last committed scrape,
// together with the input value they were built from.
//
// The zero value is ready to use. A Cache may be used by concurrent scrapes:
// each scrape reads the last committed generation, which is never modified,
// and the last scrape to commit wins.
type Cache[K, V comparable] struct {
	mu      sync.Mutex
	entries map[K]entry[V]
	// resets counts Reset calls, so that a scrape begun before a Reset
	// doesn't commit afterwards.
	resets uint64

	// scrapeMu guards active and the creation of current, so that the
	// workers of a scrape create one scrape and none outlives EndScrape.
	scrapeMu sync.Mutex
	// active is set between BeginScrape and EndScrape.
	active bool
	// current is the scrape that Scrape created after BeginScrape, or nil.
	current atomic.Pointer[Scrape[K, V]]
}

type entry[V comparable] struct {
	value   V
	metrics []prometheus.Metric
}

// Scrape collects the entries of one scrape. Its methods may be called by the
// concurrent workers of one scrape. A nil Scrape builds every metric and
// caches nothing.
type Scrape[K, V comparable] struct {
	cache  *Cache[K, V]
	resets uint64
	// prev is the last committed generation. It is never modified, so it is
	// read without holding mu.
	prev map[K]entry[V]

	mu sync.Mutex
	// next collects the entries of this scrape. Commit hands it to the cache
	// and sets it to nil, so later calls can't modify a committed generation.
	next map[K]entry[V]
}

// Begin starts a scrape based on the last committed generation.
//
// Collectors don't call Begin: the collection begins and ends their scrapes
// through BeginScrape and EndScrape, and Collect gets its scrape from
// [Cache.Scrape].
func (c *Cache[K, V]) Begin() *Scrape[K, V] {
	c.mu.Lock()
	prev, resets := c.entries, c.resets
	c.mu.Unlock()

	return &Scrape[K, V]{
		cache:  c,
		resets: resets,
		prev:   prev,
		next:   make(map[K]entry[V], len(prev)),
	}
}

// BeginScrape begins a scrape. The collection calls it before Collect.
func (c *Cache[K, V]) BeginScrape() {
	c.scrapeMu.Lock()
	c.current.Store(nil)
	c.active = true
	c.scrapeMu.Unlock()
}

// EndScrape ends the scrape begun by BeginScrape. The collection calls it
// once Collect returned err.
//
// A scrape that Collect never asked for leaves the cache unchanged. Otherwise
// the scrape is committed, unless Collect failed before it loaded or stored an
// entry: then the entries of the previous scrape stay, so a failed read of the
// objects doesn't drop them.
func (c *Cache[K, V]) EndScrape(err error) {
	c.scrapeMu.Lock()
	c.active = false
	s := c.current.Swap(nil)
	c.scrapeMu.Unlock()

	if s == nil {
		return
	}

	s.mu.Lock()
	empty := len(s.next) == 0
	s.mu.Unlock()

	if err != nil && empty {
		s.discard()

		return
	}

	s.Commit()
}

// Scrape returns the scrape of the running Collect call, which it creates on
// the first call. It is nil outside of a scrape begun by the collection, for
// example when Collect is called directly.
//
// Only Collect and the workers it waits for may call Scrape: after Collect
// returned, it returns nil or the scrape of a later Collect call.
func (c *Cache[K, V]) Scrape() *Scrape[K, V] {
	if s := c.current.Load(); s != nil {
		return s
	}

	c.scrapeMu.Lock()
	defer c.scrapeMu.Unlock()

	// Another worker of the scrape may have created it meanwhile.
	if s := c.current.Load(); s != nil {
		return s
	}

	if !c.active {
		return nil
	}

	s := c.Begin()
	c.current.Store(s)

	return s
}

// Reset drops all cached metrics, for example after their descriptors changed.
// Scrapes begun before Reset don't commit their entries.
func (c *Cache[K, V]) Reset() {
	c.mu.Lock()
	c.entries = nil
	c.resets++
	c.mu.Unlock()
}

// Load returns the metrics of the last scrape for key if they were built from
// an equal value, and keeps them for the next scrape. The caller must not
// modify the returned slice.
func (s *Scrape[K, V]) Load(key K, value V) ([]prometheus.Metric, bool) {
	if s == nil {
		return nil, false
	}

	e, ok := s.prev[key]
	if !ok || e.value != value {
		return nil, false
	}

	s.mu.Lock()
	if s.next != nil {
		s.next[key] = e
	}
	s.mu.Unlock()

	return e.metrics, true
}

// Store keeps metrics built from value for key. The caller must not modify
// metrics afterwards.
func (s *Scrape[K, V]) Store(key K, value V, metrics []prometheus.Metric) {
	if s == nil {
		return
	}

	s.mu.Lock()
	if s.next != nil {
		s.next[key] = entry[V]{value: value, metrics: metrics}
	}
	s.mu.Unlock()
}

// Get returns the metrics of the last scrape for key if they were built from
// an equal value. Otherwise it calls build and stores its result. The caller
// must not modify the returned slice.
func (s *Scrape[K, V]) Get(key K, value V, build func() []prometheus.Metric) []prometheus.Metric {
	if metrics, ok := s.Load(key, value); ok {
		return metrics
	}

	metrics := build()
	s.Store(key, value, metrics)

	return metrics
}

// Commit replaces the cached entries with those loaded or stored in this
// scrape, so keys that were not seen are dropped. It does nothing if the cache
// was reset after the scrape began, for example by a Close or Build that ran
// while a timed-out scrape was still collecting. Calls on the scrape after
// Commit don't change the cache.
func (s *Scrape[K, V]) Commit() {
	if s == nil {
		return
	}

	next := s.discard()
	if next == nil {
		// Committed before.
		return
	}

	s.cache.mu.Lock()
	if s.cache.resets == s.resets {
		s.cache.entries = next
	}
	s.cache.mu.Unlock()
}

// discard ends the scrape without changing the cache and returns its entries.
func (s *Scrape[K, V]) discard() map[K]entry[V] {
	s.mu.Lock()
	defer s.mu.Unlock()

	next := s.next
	s.next = nil

	return next
}