✨ 迁移项目

pkg/metric/counters.go

@@ -0,0 +1,117 @@
+package metric
+
+import (
+	"strings"
+	"time"
+
+	"github.com/mitchellh/mapstructure"
+	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/common/model"
+)
+
+const (
+	CounterInc = "inc"
+	CounterAdd = "add"
+
+	ErrCounterActionRequired          = "counter action must be defined as either `inc` or `add`"
+	ErrCounterInvalidAction           = "action %s is not valid, action must be either `inc` or `add`"
+	ErrCounterInvalidMatchAll         = "`match_all: true` cannot be combined with `value`, please remove `match_all` or `value`"
+	ErrCounterInvalidCountBytes       = "`count_entry_bytes: true` can only be set with `match_all: true`"
+	ErrCounterInvalidCountBytesAction = "`count_entry_bytes: true` can only be used with `action: add`"
+)
+
+type CounterConfig struct {
+	MatchAll   *bool   `mapstructure:"match_all"`
+	CountBytes *bool   `mapstructure:"count_entry_bytes"`
+	Value      *string `mapstructure:"value"`
+	Action     string  `mapstructure:"action"`
+}
+
+func validateCounterConfig(config *CounterConfig) error {
+	if config.Action == "" {
+		return errors.New(ErrCounterActionRequired)
+	}
+	config.Action = strings.ToLower(config.Action)
+	if config.Action != CounterInc && config.Action != CounterAdd {
+		return errors.Errorf(ErrCounterInvalidAction, config.Action)
+	}
+	if config.MatchAll != nil && *config.MatchAll && config.Value != nil {
+		return errors.Errorf(ErrCounterInvalidMatchAll)
+	}
+	if config.CountBytes != nil && *config.CountBytes && (config.MatchAll == nil || !*config.MatchAll) {
+		return errors.New(ErrCounterInvalidCountBytes)
+	}
+	if config.CountBytes != nil && *config.CountBytes && config.Action != CounterAdd {
+		return errors.New(ErrCounterInvalidCountBytesAction)
+	}
+	return nil
+}
+
+func parseCounterConfig(config interface{}) (*CounterConfig, error) {
+	cfg := &CounterConfig{}
+	err := mapstructure.Decode(config, cfg)
+	if err != nil {
+		return nil, err
+	}
+	return cfg, nil
+}
+
+// Counters is a vec tor of counters for a each log stream.
+type Counters struct {
+	*metricVec
+	Cfg *CounterConfig
+}
+
+// NewCounters creates a new counter vec.
+func NewCounters(name, help string, config interface{}, maxIdleSec int64) (*Counters, error) {
+	cfg, err := parseCounterConfig(config)
+	if err != nil {
+		return nil, err
+	}
+	err = validateCounterConfig(cfg)
+	if err != nil {
+		return nil, err
+	}
+	return &Counters{
+		metricVec: newMetricVec(func(labels map[string]string) prometheus.Metric {
+			return &expiringCounter{prometheus.NewCounter(prometheus.CounterOpts{
+				Help:        help,
+				Name:        name,
+				ConstLabels: labels,
+			}),
+				0,
+			}
+		}, maxIdleSec),
+		Cfg: cfg,
+	}, nil
+}
+
+// With returns the counter associated with a stream labelset.
+func (c *Counters) With(labels model.LabelSet) prometheus.Counter {
+	return c.metricVec.With(labels).(prometheus.Counter)
+}
+
+type expiringCounter struct {
+	prometheus.Counter
+	lastModSec int64
+}
+
+// Inc increments the counter by 1. Use Add to increment it by arbitrary
+// non-negative values.
+func (e *expiringCounter) Inc() {
+	e.Counter.Inc()
+	e.lastModSec = time.Now().Unix()
+}
+
+// Add adds the given value to the counter. It panics if the value is <
+// 0.
+func (e *expiringCounter) Add(val float64) {
+	e.Counter.Add(val)
+	e.lastModSec = time.Now().Unix()
+}
+
+// HasExpired implements Expirable
+func (e *expiringCounter) HasExpired(currentTimeSec int64, maxAgeSec int64) bool {
+	return currentTimeSec-e.lastModSec >= maxAgeSec
+}

pkg/metric/counters_test.go

@@ -0,0 +1,140 @@
+package metric
+
+import (
+	"testing"
+	"time"
+
+	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/common/model"
+	"github.com/stretchr/testify/assert"
+)
+
+var (
+	counterTestTrue  = true
+	counterTestFalse = false
+	counterTestVal   = "some val"
+)
+
+func Test_validateCounterConfig(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name   string
+		config CounterConfig
+		err    error
+	}{
+		{"invalid action",
+			CounterConfig{
+				Action: "del",
+			},
+			errors.Errorf(ErrCounterInvalidAction, "del"),
+		},
+		{"invalid counter match all",
+			CounterConfig{
+				MatchAll: &counterTestTrue,
+				Value:    &counterTestVal,
+				Action:   "inc",
+			},
+			errors.New(ErrCounterInvalidMatchAll),
+		},
+		{"invalid counter match bytes",
+			CounterConfig{
+				MatchAll:   nil,
+				CountBytes: &counterTestTrue,
+				Action:     "add",
+			},
+			errors.New(ErrCounterInvalidCountBytes),
+		},
+		{"invalid counter match bytes action",
+			CounterConfig{
+				MatchAll:   &counterTestTrue,
+				CountBytes: &counterTestTrue,
+				Action:     "inc",
+			},
+			errors.New(ErrCounterInvalidCountBytesAction),
+		},
+		{"valid counter match bytes",
+			CounterConfig{
+				MatchAll:   &counterTestTrue,
+				CountBytes: &counterTestTrue,
+				Action:     "add",
+			},
+			nil,
+		},
+		{"valid",
+			CounterConfig{
+				Value:  &counterTestVal,
+				Action: "inc",
+			},
+			nil,
+		},
+		{"valid match all is false",
+			CounterConfig{
+				MatchAll: &counterTestFalse,
+				Value:    &counterTestVal,
+				Action:   "inc",
+			},
+			nil,
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			err := validateCounterConfig(&tt.config)
+			if ((err != nil) && (err.Error() != tt.err.Error())) || (err == nil && tt.err != nil) {
+				t.Errorf("Metrics stage validation error, expected error = %v, actual error = %v", tt.err, err)
+				return
+			}
+		})
+	}
+}
+
+func TestCounterExpiration(t *testing.T) {
+	t.Parallel()
+	cfg := CounterConfig{
+		Action: "inc",
+	}
+
+	cnt, err := NewCounters("test1", "HELP ME!!!!!", cfg, 1)
+	assert.Nil(t, err)
+
+	// Create a label and increment the counter
+	lbl1 := model.LabelSet{}
+	lbl1["test"] = "i don't wanna make this a constant"
+	cnt.With(lbl1).Inc()
+
+	// Collect the metrics, should still find the metric in the map
+	collect(cnt)
+	assert.Contains(t, cnt.metrics, lbl1.Fingerprint())
+
+	time.Sleep(1100 * time.Millisecond) // Wait just past our max idle of 1 sec
+
+	//Add another counter with new label val
+	lbl2 := model.LabelSet{}
+	lbl2["test"] = "eat this linter"
+	cnt.With(lbl2).Inc()
+
+	// Collect the metrics, first counter should have expired and removed, second should still be present
+	collect(cnt)
+	assert.NotContains(t, cnt.metrics, lbl1.Fingerprint())
+	assert.Contains(t, cnt.metrics, lbl2.Fingerprint())
+}
+
+func collect(c prometheus.Collector) {
+	done := make(chan struct{})
+	collector := make(chan prometheus.Metric)
+
+	go func() {
+		defer close(done)
+		c.Collect(collector)
+	}()
+
+	for {
+		select {
+		case <-collector:
+		case <-done:
+			return
+		}
+	}
+}

pkg/metric/gauges.go

@@ -0,0 +1,136 @@
+package metric
+
+import (
+	"strings"
+	"time"
+
+	"github.com/mitchellh/mapstructure"
+	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/common/model"
+)
+
+const (
+	GaugeSet = "set"
+	GaugeInc = "inc"
+	GaugeDec = "dec"
+	GaugeAdd = "add"
+	GaugeSub = "sub"
+
+	ErrGaugeActionRequired = "gauge action must be defined as `set`, `inc`, `dec`, `add`, or `sub`"
+	ErrGaugeInvalidAction  = "action %s is not valid, action must be `set`, `inc`, `dec`, `add`, or `sub`"
+)
+
+type GaugeConfig struct {
+	Value  *string `mapstructure:"value"`
+	Action string  `mapstructure:"action"`
+}
+
+func validateGaugeConfig(config *GaugeConfig) error {
+	if config.Action == "" {
+		return errors.New(ErrGaugeActionRequired)
+	}
+	config.Action = strings.ToLower(config.Action)
+	if config.Action != GaugeSet &&
+		config.Action != GaugeInc &&
+		config.Action != GaugeDec &&
+		config.Action != GaugeAdd &&
+		config.Action != GaugeSub {
+		return errors.Errorf(ErrGaugeInvalidAction, config.Action)
+	}
+	return nil
+}
+
+func parseGaugeConfig(config interface{}) (*GaugeConfig, error) {
+	cfg := &GaugeConfig{}
+	err := mapstructure.Decode(config, cfg)
+	if err != nil {
+		return nil, err
+	}
+	return cfg, nil
+}
+
+// Gauges is a vector of gauges for a each log stream.
+type Gauges struct {
+	*metricVec
+	Cfg *GaugeConfig
+}
+
+// NewGauges creates a new gauge vec.
+func NewGauges(name, help string, config interface{}, maxIdleSec int64) (*Gauges, error) {
+	cfg, err := parseGaugeConfig(config)
+	if err != nil {
+		return nil, err
+	}
+	err = validateGaugeConfig(cfg)
+	if err != nil {
+		return nil, err
+	}
+	return &Gauges{
+		metricVec: newMetricVec(func(labels map[string]string) prometheus.Metric {
+			return &expiringGauge{prometheus.NewGauge(prometheus.GaugeOpts{
+				Help:        help,
+				Name:        name,
+				ConstLabels: labels,
+			}),
+				0,
+			}
+		}, maxIdleSec),
+		Cfg: cfg,
+	}, nil
+}
+
+// With returns the gauge associated with a stream labelset.
+func (g *Gauges) With(labels model.LabelSet) prometheus.Gauge {
+	return g.metricVec.With(labels).(prometheus.Gauge)
+}
+
+type expiringGauge struct {
+	prometheus.Gauge
+	lastModSec int64
+}
+
+// Set sets the Gauge to an arbitrary value.
+func (g *expiringGauge) Set(val float64) {
+	g.Gauge.Set(val)
+	g.lastModSec = time.Now().Unix()
+}
+
+// Inc increments the Gauge by 1. Use Add to increment it by arbitrary
+// values.
+func (g *expiringGauge) Inc() {
+	g.Gauge.Inc()
+	g.lastModSec = time.Now().Unix()
+}
+
+// Dec decrements the Gauge by 1. Use Sub to decrement it by arbitrary
+// values.
+func (g *expiringGauge) Dec() {
+	g.Gauge.Dec()
+	g.lastModSec = time.Now().Unix()
+}
+
+// Add adds the given value to the Gauge. (The value can be negative,
+// resulting in a decrease of the Gauge.)
+func (g *expiringGauge) Add(val float64) {
+	g.Gauge.Add(val)
+	g.lastModSec = time.Now().Unix()
+}
+
+// Sub subtracts the given value from the Gauge. (The value can be
+// negative, resulting in an increase of the Gauge.)
+func (g *expiringGauge) Sub(val float64) {
+	g.Gauge.Sub(val)
+	g.lastModSec = time.Now().Unix()
+}
+
+// SetToCurrentTime sets the Gauge to the current Unix time in seconds.
+func (g *expiringGauge) SetToCurrentTime() {
+	g.Gauge.SetToCurrentTime()
+	g.lastModSec = time.Now().Unix()
+}
+
+// HasExpired implements Expirable
+func (g *expiringGauge) HasExpired(currentTimeSec int64, maxAgeSec int64) bool {
+	return currentTimeSec-g.lastModSec >= maxAgeSec
+}

pkg/metric/gauges_test.go

@@ -0,0 +1,40 @@
+package metric
+
+import (
+	"testing"
+	"time"
+
+	"github.com/prometheus/common/model"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGaugeExpiration(t *testing.T) {
+	t.Parallel()
+	cfg := GaugeConfig{
+		Action: "inc",
+	}
+
+	gag, err := NewGauges("test1", "HELP ME!!!!!", cfg, 1)
+	assert.Nil(t, err)
+
+	// Create a label and increment the gauge
+	lbl1 := model.LabelSet{}
+	lbl1["test"] = "app"
+	gag.With(lbl1).Inc()
+
+	// Collect the metrics, should still find the metric in the map
+	collect(gag)
+	assert.Contains(t, gag.metrics, lbl1.Fingerprint())
+
+	time.Sleep(1100 * time.Millisecond) // Wait just past our max idle of 1 sec
+
+	//Add another gauge with new label val
+	lbl2 := model.LabelSet{}
+	lbl2["test"] = "app2"
+	gag.With(lbl2).Inc()
+
+	// Collect the metrics, first gauge should have expired and removed, second should still be present
+	collect(gag)
+	assert.NotContains(t, gag.metrics, lbl1.Fingerprint())
+	assert.Contains(t, gag.metrics, lbl2.Fingerprint())
+}

pkg/metric/histograms.go

@@ -0,0 +1,79 @@
+package metric
+
+import (
+	"time"
+
+	"github.com/mitchellh/mapstructure"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/common/model"
+)
+
+type HistogramConfig struct {
+	Value   *string   `mapstructure:"value"`
+	Buckets []float64 `mapstructure:"buckets"`
+}
+
+func validateHistogramConfig(config *HistogramConfig) error {
+	return nil
+}
+
+func parseHistogramConfig(config interface{}) (*HistogramConfig, error) {
+	cfg := &HistogramConfig{}
+	err := mapstructure.Decode(config, cfg)
+	if err != nil {
+		return nil, err
+	}
+	return cfg, nil
+}
+
+// Histograms is a vector of histograms for a each log stream.
+type Histograms struct {
+	*metricVec
+	Cfg *HistogramConfig
+}
+
+// NewHistograms creates a new histogram vec.
+func NewHistograms(name, help string, config interface{}, maxIdleSec int64) (*Histograms, error) {
+	cfg, err := parseHistogramConfig(config)
+	if err != nil {
+		return nil, err
+	}
+	err = validateHistogramConfig(cfg)
+	if err != nil {
+		return nil, err
+	}
+	return &Histograms{
+		metricVec: newMetricVec(func(labels map[string]string) prometheus.Metric {
+			return &expiringHistogram{prometheus.NewHistogram(prometheus.HistogramOpts{
+				Help:        help,
+				Name:        name,
+				ConstLabels: labels,
+				Buckets:     cfg.Buckets,
+			}),
+				0,
+			}
+		}, maxIdleSec),
+		Cfg: cfg,
+	}, nil
+}
+
+// With returns the histogram associated with a stream labelset.
+func (h *Histograms) With(labels model.LabelSet) prometheus.Histogram {
+	return h.metricVec.With(labels).(prometheus.Histogram)
+}
+
+type expiringHistogram struct {
+	prometheus.Histogram
+	lastModSec int64
+}
+
+// Observe adds a single observation to the histogram.
+func (h *expiringHistogram) Observe(val float64) {
+	h.Histogram.Observe(val)
+	h.lastModSec = time.Now().Unix()
+}
+
+// HasExpired implements Expirable
+func (h *expiringHistogram) HasExpired(currentTimeSec int64, maxAgeSec int64) bool {
+	return currentTimeSec-h.lastModSec >= maxAgeSec
+}

pkg/metric/histograms_test.go

@@ -0,0 +1,38 @@
+package metric
+
+import (
+	"testing"
+	"time"
+
+	"github.com/prometheus/common/model"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHistogramExpiration(t *testing.T) {
+	t.Parallel()
+	cfg := HistogramConfig{}
+
+	hist, err := NewHistograms("test1", "HELP ME!!!!!", cfg, 1)
+	assert.Nil(t, err)
+
+	// Create a label and increment the histogram
+	lbl1 := model.LabelSet{}
+	lbl1["test"] = "app"
+	hist.With(lbl1).Observe(23)
+
+	// Collect the metrics, should still find the metric in the map
+	collect(hist)
+	assert.Contains(t, hist.metrics, lbl1.Fingerprint())
+
+	time.Sleep(1100 * time.Millisecond) // Wait just past our max idle of 1 sec
+
+	//Add another histogram with new label val
+	lbl2 := model.LabelSet{}
+	lbl2["test"] = "app2"
+	hist.With(lbl2).Observe(2)
+
+	// Collect the metrics, first histogram should have expired and removed, second should still be present
+	collect(hist)
+	assert.NotContains(t, hist.metrics, lbl1.Fingerprint())
+	assert.Contains(t, hist.metrics, lbl2.Fingerprint())
+}

pkg/metric/metricvec.go

@@ -0,0 +1,82 @@
+package metric
+
+import (
+	"sync"
+	"time"
+
+	"github.com/lixh00/loki-client-go/pkg/labelutil"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/common/model"
+)
+
+// Expirable allows checking if something has exceeded the provided maxAge based on the provided currentTime
+type Expirable interface {
+	HasExpired(currentTimeSec int64, maxAgeSec int64) bool
+}
+
+type metricVec struct {
+	factory   func(labels map[string]string) prometheus.Metric
+	mtx       sync.Mutex
+	metrics   map[model.Fingerprint]prometheus.Metric
+	maxAgeSec int64
+}
+
+func newMetricVec(factory func(labels map[string]string) prometheus.Metric, maxAgeSec int64) *metricVec {
+	return &metricVec{
+		metrics:   map[model.Fingerprint]prometheus.Metric{},
+		factory:   factory,
+		maxAgeSec: maxAgeSec,
+	}
+}
+
+// Describe implements prometheus.Collector and doesn't declare any metrics on purpose to bypass prometheus validation.
+// see https://godoc.org/github.com/prometheus/client_golang/prometheus#hdr-Custom_Collectors_and_constant_Metrics search for "unchecked"
+func (c *metricVec) Describe(ch chan<- *prometheus.Desc) {}
+
+// Collect implements prometheus.Collector
+func (c *metricVec) Collect(ch chan<- prometheus.Metric) {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+	for _, m := range c.metrics {
+		ch <- m
+	}
+	c.prune()
+}
+
+// With returns the metric associated with the labelset.
+func (c *metricVec) With(labels model.LabelSet) prometheus.Metric {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+	fp := labels.Fingerprint()
+	var ok bool
+	var metric prometheus.Metric
+	if metric, ok = c.metrics[fp]; !ok {
+		metric = c.factory(labelutil.ModelLabelSetToMap(labels))
+		c.metrics[fp] = metric
+	}
+	return metric
+}
+
+func (c *metricVec) Delete(labels model.LabelSet) bool {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+	fp := labels.Fingerprint()
+	_, ok := c.metrics[fp]
+	if ok {
+		delete(c.metrics, fp)
+	}
+	return ok
+}
+
+// prune will remove all metrics which implement the Expirable interface and have expired
+// it does not take out a lock on the metrics map so whoever calls this function should do so.
+func (c *metricVec) prune() {
+	currentTimeSec := time.Now().Unix()
+	for fp, m := range c.metrics {
+		if em, ok := m.(Expirable); ok {
+			if em.HasExpired(currentTimeSec, c.maxAgeSec) {
+				delete(c.metrics, fp)
+			}
+		}
+	}
+}