xk6-frostfs/internal/stats/stats.go

package stats

import (
	"strings"
	"time"

	"go.k6.io/k6/js/modules"
	"go.k6.io/k6/metrics"
)

// RootModule is the global module object type. It is instantiated once per test
// run and will be used to create k6/x/frostfs/stats module instances for each VU.
type RootModule struct {
	Instance string
}

var (
	tagSet *metrics.TagSet

	Registry *metrics.Registry
)

func init() {
	Registry = metrics.NewRegistry()
	tagSet = Registry.RootTagSet()
	modules.Register("k6/x/frostfs/stats", &RootModule{})
}

// SetTags sets additional tags to custom metrics.
// Format: "key1:value1;key2:value2".
// Panics if input has invalid format.
func (m *RootModule) SetTags(labels string) {
	kv := make(map[string]string)
	pairs := strings.Split(labels, ";")
	for _, pair := range pairs {
		items := strings.Split(pair, ":")
		if len(items) != 2 {
			panic("invalid labels format")
		}
		kv[strings.TrimSpace(items[0])] = strings.TrimSpace(items[1])
	}
	for k, v := range kv {
		tagSet = tagSet.With(k, v)
	}
}

func Report(vu modules.VU, metric *metrics.Metric, value float64) {
	metrics.PushIfNotDone(vu.Context(), vu.State().Samples, metrics.Sample{
		TimeSeries: metrics.TimeSeries{
			Metric: metric,
			Tags:   tagSet,
		},
		Time:  time.Now(),
		Value: value,
	})
}

func ReportDataReceived(vu modules.VU, value float64) {
	vu.State().BuiltinMetrics.DataReceived.Sink.Add(
		metrics.Sample{
			TimeSeries: metrics.TimeSeries{
				Metric: &metrics.Metric{},
				Tags:   tagSet,
			},
			Value: value,
			Time:  time.Now(),
		},
	)
}

func ReportDataSent(vu modules.VU, value float64) {
	state := vu.State()
	state.BuiltinMetrics.DataSent.Sink.Add(
		metrics.Sample{
			TimeSeries: metrics.TimeSeries{
				Metric: &metrics.Metric{},
				Tags:   tagSet,
			},
			Value: value,
			Time:  time.Now(),
		},
	)
}