clustvirt/lib/stats/stats.go

package stats

import (
	"context"
	"errors"
	"fmt"
	"io"
	"time"

	"golang.org/x/exp/constraints"
)

type Number interface {
	constraints.Integer | constraints.Float
}

// StatFetcher represents a data collector. It is given a channel to send a Number
// to, being the metric that is being collected.
type StatFetcher[T Number] func(chan<- Stat[T])

// StatFetcherList is a map of collector types, where the key is the label to use
// for the statistic
type StatFetcherList[T Number] map[string]StatFetcher[T]

// StatProvider is something that has the Stats method and provides stats
type StatProvider[T Number] interface {
	Stats() StatFetcherList[T]
}

// StatCollector is something that has the Collect method, recieving a StatProvider to collect from
type StatCollector[T Number] interface {
	Collect(StatProvider[T])
}

// Stat represents a single statistic
type Stat[T Number] struct {
	Name  string
	Value T
}

// StatsCollector controls running stats
type StatsCollector[T Number] struct {
	ctx        context.Context
	ctxCancel  context.CancelFunc
	collectors StatFetcherList[T]
	interval   time.Duration
	rchan      chan Stat[T]
	running    bool
}

// AddCollector adds a collector to the existing list of collectors. If `name` is already defined,
// will return an error but still replace the collector. If the StatsCollector is running, it is restarted.
func (sc *StatsCollector[T]) AddCollector(name string, collector StatFetcher[T], restart bool) error {
	var err error
	if _, ok := sc.collectors[name]; !ok {
		err = fmt.Errorf("replacing collector %s", name)
	}
	sc.collectors[name] = collector
	if sc.IsRunning() {
		sc.Restart()
	}
	return err
}

// MergeCollectors is a utility function to merge a Collectors map with another one, allowing for
// better localization of code and avoiding repeating the same function all over the place when
// an object with Collector functions wants to add them
func (sc *StatsCollector[T]) MergeCollectors(collectors StatFetcherList[T]) error {
	var errList []error
	for k, c := range collectors {
		errList = append(errList, sc.AddCollector(k, c, false))
	}
	sc.Restart()
	return errors.Join(errList...)
}

// Collect satisfys the StatCollector interface. It is given a StatProvider, and those stats are then
// merged with the existing providers, running on the same interval
func (sc *StatsCollector[T]) Collect(sp StatProvider[T]) {
	sc.MergeCollectors(sp.Stats())
}

// IsRunning is a utility function that returns true if the StatsCollector is running
func (sc *StatsCollector[T]) IsRunning() bool { return sc.running }

// RunningStats returns a channel of type Stat, and receives a context and a list
// of collectors that get run every 'd' duration
func (sc *StatsCollector[T]) Start() {
	sc.ctx, sc.ctxCancel = context.WithCancel(context.Background())
	sc.rchan = make(chan Stat[T])
	ticker := time.NewTicker(sc.interval)

	// start the goroutine that does the actual work
	go func() {
		defer func() { sc.running = false }()
		select {
		case <-sc.ctx.Done():
			return
		case <-ticker.C:
			for _, f := range sc.collectors {
				f(sc.rchan)
			}
		}
	}()
	sc.running = true
}

// Stop stops the collection of stats
func (sc *StatsCollector[T]) Stop() {
	defer close(sc.rchan)
	sc.ctxCancel()
	<-sc.ctx.Done()
}

// Restart stops the StatsCollector,
func (sc *StatsCollector[T]) Restart() {
	if sc.IsRunning() {
		sc.Stop()
	}
	if !sc.IsRunning() {
		sc.Start()
	}
}

// StatList is a simple map[string]Stat tht allows for easy lookup of stats by name
// Getting all stats can still be done using a range loop
type StatList[T Number] map[string]Stat[T]

func (sl StatList[T]) Get(name string) Stat[T] {
	if v, ok := sl[name]; ok {
		return v
	}
	return Stat[T]{
		Name: "UNDEFINED",
		Value: 0,
	}
}

func (sl StatList[T]) Set(name string, s Stat[T]) {
	sl[name] = s
}

// PrometheusMetrics writes the stats in a format prometheus understands to an io.Writer
func (sl StatList[T]) PrometheusMetrics(w io.Writer) {

}