service.go

package labelstore

import (
	"context"
	"sync"
	"time"

	"github.com/go-kit/log"
	"github.com/grafana/agent/pkg/flow/logging/level"
	agent_service "github.com/grafana/agent/service"
	flow_service "github.com/grafana/agent/service"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/prometheus/model/labels"
)

const ServiceName = "labelstore"

type service struct {
	log                 log.Logger
	mut                 sync.Mutex
	globalRefID         uint64
	mappings            map[string]*remoteWriteMapping
	labelsHashToGlobal  map[uint64]uint64
	staleGlobals        map[uint64]*staleMarker
	totalIDs            *prometheus.Desc
	idsInRemoteWrapping *prometheus.Desc
	lastStaleCheck      prometheus.Gauge
}
type staleMarker struct {
	globalID        uint64
	lastMarkedStale time.Time
	labelHash       uint64
}

type Arguments struct{}

var _ flow_service.Service = (*service)(nil)

func New(l log.Logger, r prometheus.Registerer) *service {
	if l == nil {
		l = log.NewNopLogger()
	}
	s := &service{
		log:                 l,
		globalRefID:         0,
		mappings:            make(map[string]*remoteWriteMapping),
		labelsHashToGlobal:  make(map[uint64]uint64),
		staleGlobals:        make(map[uint64]*staleMarker),
		totalIDs:            prometheus.NewDesc("agent_labelstore_global_ids_count", "Total number of global ids.", nil, nil),
		idsInRemoteWrapping: prometheus.NewDesc("agent_labelstore_remote_store_ids_count", "Total number of ids per remote write", []string{"remote_name"}, nil),
		lastStaleCheck: prometheus.NewGauge(prometheus.GaugeOpts{
			Name: "agent_labelstore_last_stale_check_timestamp",
			Help: "Last time stale check was ran expressed in unix timestamp.",
		}),
	}
	_ = r.Register(s.lastStaleCheck)
	_ = r.Register(s)
	return s
}

// Definition returns the Definition of the Service.
// Definition must always return the same value across all
// calls.
func (s *service) Definition() agent_service.Definition {
	return agent_service.Definition{
		Name:       ServiceName,
		ConfigType: Arguments{},
		DependsOn:  nil,
	}
}

func (s *service) Describe(m chan<- *prometheus.Desc) {
	m <- s.totalIDs
	m <- s.idsInRemoteWrapping
}
func (s *service) Collect(m chan<- prometheus.Metric) {
	s.mut.Lock()
	defer s.mut.Unlock()

	m <- prometheus.MustNewConstMetric(s.totalIDs, prometheus.GaugeValue, float64(len(s.labelsHashToGlobal)))
	for name, rw := range s.mappings {
		m <- prometheus.MustNewConstMetric(s.idsInRemoteWrapping, prometheus.GaugeValue, float64(len(rw.globalToLocal)), name)
	}
}

// Run starts a Service. Run must block until the provided
// context is canceled. Returning an error should be treated
// as a fatal error for the Service.
func (s *service) Run(ctx context.Context, host agent_service.Host) error {
	staleCheck := time.NewTicker(10 * time.Minute)
	for {
		select {
		case <-ctx.Done():
			return nil
		case <-staleCheck.C:
			s.CheckAndRemoveStaleMarkers()
		}
	}
}

// Update updates a Service at runtime. Update is never
// called if [Definition.ConfigType] is nil. newConfig will
// be the same type as ConfigType; if ConfigType is a
// pointer to a type, newConfig will be a pointer to the
// same type.
//
// Update will be called once before Run, and may be called
// while Run is active.
func (s *service) Update(_ any) error {
	return nil
}

// Data returns the Data associated with a Service. Data
// must always return the same value across multiple calls,
// as callers are expected to be able to cache the result.
//
// Data may be invoked before Run.
func (s *service) Data() any {
	return s
}

// GetOrAddLink is called by a remote_write endpoint component to add mapping and get back the global id.
func (s *service) GetOrAddLink(componentID string, localRefID uint64, lbls labels.Labels) uint64 {
	s.mut.Lock()
	defer s.mut.Unlock()

	// If the mapping doesn't exist then we need to create it
	m, found := s.mappings[componentID]
	if !found {
		m = &remoteWriteMapping{
			RemoteWriteID: componentID,
			localToGlobal: make(map[uint64]uint64),
			globalToLocal: make(map[uint64]uint64),
		}
		s.mappings[componentID] = m
	}

	labelHash := lbls.Hash()
	globalID, found := s.labelsHashToGlobal[labelHash]
	if found {
		m.localToGlobal[localRefID] = globalID
		m.globalToLocal[globalID] = localRefID
		return globalID
	}
	// We have a value we have never seen before so increment the globalrefid and assign
	s.globalRefID++
	s.labelsHashToGlobal[labelHash] = s.globalRefID
	m.localToGlobal[localRefID] = s.globalRefID
	m.globalToLocal[s.globalRefID] = localRefID
	return s.globalRefID
}

// GetOrAddGlobalRefID is used to create a global refid for a labelset
func (s *service) GetOrAddGlobalRefID(l labels.Labels) uint64 {
	s.mut.Lock()
	defer s.mut.Unlock()

	// Guard against bad input.
	if l == nil {
		return 0
	}

	labelHash := l.Hash()
	globalID, found := s.labelsHashToGlobal[labelHash]
	if found {
		return globalID
	}
	s.globalRefID++
	s.labelsHashToGlobal[labelHash] = s.globalRefID
	return s.globalRefID
}

// GetGlobalRefID returns the global refid for a component local combo, or 0 if not found
func (s *service) GetGlobalRefID(componentID string, localRefID uint64) uint64 {
	s.mut.Lock()
	defer s.mut.Unlock()

	m, found := s.mappings[componentID]
	if !found {
		return 0
	}
	global := m.localToGlobal[localRefID]
	return global
}

// GetLocalRefID returns the local refid for a component global combo, or 0 if not found
func (s *service) GetLocalRefID(componentID string, globalRefID uint64) uint64 {
	s.mut.Lock()
	defer s.mut.Unlock()

	m, found := s.mappings[componentID]
	if !found {
		return 0
	}
	local := m.globalToLocal[globalRefID]
	return local
}

// AddStaleMarker adds a stale marker
func (s *service) AddStaleMarker(globalRefID uint64, l labels.Labels) {
	s.mut.Lock()
	defer s.mut.Unlock()

	s.staleGlobals[globalRefID] = &staleMarker{
		lastMarkedStale: time.Now(),
		labelHash:       l.Hash(),
		globalID:        globalRefID,
	}
}

// RemoveStaleMarker removes a stale marker
func (s *service) RemoveStaleMarker(globalRefID uint64) {
	s.mut.Lock()
	defer s.mut.Unlock()

	delete(s.staleGlobals, globalRefID)
}

// staleDuration determines how long we should wait after a stale value is received to GC that value
var staleDuration = time.Minute * 10

// CheckAndRemoveStaleMarkers is called to garbage collect and items that have grown stale over stale duration (10m)
func (s *service) CheckAndRemoveStaleMarkers() {
	s.mut.Lock()
	defer s.mut.Unlock()

	s.lastStaleCheck.Set(float64(time.Now().Unix()))
	level.Debug(s.log).Log("msg", "labelstore removing stale markers")
	curr := time.Now()
	idsToBeGCed := make([]*staleMarker, 0)
	for _, stale := range s.staleGlobals {
		// If the difference between now and the last time the stale was marked doesn't exceed stale then let it stay
		if curr.Sub(stale.lastMarkedStale) < staleDuration {
			continue
		}
		idsToBeGCed = append(idsToBeGCed, stale)
	}

	level.Debug(s.log).Log("msg", "number of ids to remove", "count", len(idsToBeGCed))

	for _, marker := range idsToBeGCed {
		delete(s.staleGlobals, marker.globalID)
		delete(s.labelsHashToGlobal, marker.labelHash)
		// Delete our mapping keys
		for _, mapping := range s.mappings {
			mapping.deleteStaleIDs(marker.globalID)
		}
	}
}

func (rw *remoteWriteMapping) deleteStaleIDs(globalID uint64) {
	localID, found := rw.globalToLocal[globalID]
	if !found {
		return
	}
	delete(rw.globalToLocal, globalID)
	delete(rw.localToGlobal, localID)
}

// remoteWriteMapping maps a remote_write to a set of global ids
type remoteWriteMapping struct {
	RemoteWriteID string
	localToGlobal map[uint64]uint64
	globalToLocal map[uint64]uint64
}