Add count operation (useful if multiple metrics match the labels) (#45)

Signed-off-by: Jirka Kremser <jiri.kremser@gmail.com>

Author: jkremser

metric/mem_store.go

@@ -40,12 +40,12 @@ func (m ms) Get(unescapedName types.MetricName, searchLabels types.Labels, timeO
 	return value, found, err
 }
 
-func (m ms) get(name types.MetricName, searchLabels types.Labels, timeOp types.OperationOverTime, defaultAggregation types.AggregationOverVectors) (float64, types.Found, error) {
+func (m ms) get(name types.MetricName, searchLabels types.Labels, timeOp types.OperationOverTime, vecOp types.AggregationOverVectors) (float64, types.Found, error) {
 	now := time.Now().Unix()
 	if err := util.CheckTimeOp(timeOp); err != nil {
 		return -1., false, err
 	}
-	if err := checkDefaultAggregation(defaultAggregation); err != nil {
+	if err := checkVectorAggregation(vecOp); err != nil {
 		return -1., false, err
 	}
 	storedMetrics, found := m.store.Load(string(name))
@@ -55,6 +55,9 @@ func (m ms) get(name types.MetricName, searchLabels types.Labels, timeOp types.O
 	}
 	if md, f := storedMetrics.Load(hashOfMap(searchLabels)); f {
 		// found exact label match
+		if vecOp == types.VecCount {
+			return 1, true, nil
+		}
 		if !m.isStale(md.LastUpdate, now) {
 			ret, f := md.AggregatesOverTime.Load(timeOp)
 			if !f {
@@ -87,7 +90,7 @@ func (m ms) get(name types.MetricName, searchLabels types.Labels, timeOp types.O
 					return true
 				}
 				counter += 1
-				accumulator = m.calculateAggregate(val, counter, accumulator, defaultAggregation)
+				accumulator = m.calculateAggregate(val, counter, accumulator, vecOp)
 			} else {
 				defer func() {
 					storedMetrics.Delete(hashOfMap(searchLabels))
@@ -99,9 +102,9 @@ func (m ms) get(name types.MetricName, searchLabels types.Labels, timeOp types.O
 	return accumulator, true, nil
 }
 
-func checkDefaultAggregation(aggregation types.AggregationOverVectors) error {
+func checkVectorAggregation(aggregation types.AggregationOverVectors) error {
 	switch aggregation {
-	case types.VecSum, types.VecAvg, types.VecMin, types.VecMax:
+	case types.VecSum, types.VecAvg, types.VecMin, types.VecMax, types.VecCount:
 		return nil
 	default:
 		return fmt.Errorf("unknown AggregationOverVectors:%s", aggregation)
@@ -173,7 +176,11 @@ func (m ms) isStale(datapoint uint32, now int64) bool {
 
 func (m ms) calculateAggregate(value float64, counter int, accumulator float64, aggregation types.AggregationOverVectors) float64 {
 	if counter == 1 {
-		return value
+		if aggregation == types.VecCount {
+			return 1
+		} else {
+			return value
+		}
 	}
 	switch aggregation {
 	case types.VecSum:
@@ -188,6 +195,8 @@ func (m ms) calculateAggregate(value float64, counter int, accumulator float64,
 		return math.Min(accumulator, value)
 	case types.VecMax:
 		return math.Max(accumulator, value)
+	case types.VecCount:
+		return accumulator + 1
 	default:
 		panic("unknown aggregation function: " + aggregation)
 	}

metric/mem_store_test.go

@@ -521,6 +521,49 @@ func TestMemStoreSumOverAverages(t *testing.T) {
 	assertMetricFound(t, val, found, err, 3.5+2.333)
 }
 
+func TestMemStoreCount(t *testing.T) {
+	// setup
+	ms := NewMetricStore(60)
+	labels1 := map[string]any{
+		"a": "1",
+		"b": "2",
+	}
+	labels2 := map[string]any{
+		"a": "1",
+		"b": "3",
+	}
+	labels3 := map[string]any{
+		"a": "1",
+		"b": "4",
+	}
+	labels4 := map[string]any{
+		"a": "2",
+		"b": "2",
+	}
+	name1 := "metric_name"
+	setupMetrics(ms, name1, 1, labels1, 1., 2.)
+	setupMetrics(ms, name1, 1, labels2, 1., 2., 3.)
+	setupMetrics(ms, name1, 1, labels3, 1., 2., 3., 4.)
+	setupMetrics(ms, name1, 1, labels4, 1., 2., 3., 4., 5.)
+	setupMetrics(ms, "noise", 1, labels2, 1., 2., 3., 4., 5.) // this shouldn't be included
+	val1, found1, err1 := ms.Get(types.MetricName(name1), map[string]any{
+		"a": "1",
+	}, types.OpAvg, types.VecCount)
+	assertMetricFound(t, val1, found1, err1, 3.)
+
+	val2, found2, err2 := ms.Get(types.MetricName(name1), map[string]any{
+		"b": "2",
+	}, types.OpAvg, types.VecCount)
+	assertMetricFound(t, val2, found2, err2, 2.)
+	val3, found3, err3 := ms.Get(types.MetricName(name1), map[string]any{}, types.OpAvg, types.VecCount)
+	assertMetricFound(t, val3, found3, err3, 4.)
+	val4, found4, err4 := ms.Get(types.MetricName(name1), map[string]any{
+		"a": "1",
+		"b": "2",
+	}, types.OpAvg, types.VecCount)
+	assertMetricFound(t, val4, found4, err4, 1.)
+}
+
 func setupMetrics(store types.MemStore, name string, secondsStep int64, labels map[string]any, vals ...float64) {
 	now := time.Now().Unix()
 	for i, v := range vals {

metric/simple_parser.go

@@ -23,7 +23,7 @@ func (p p) Parse(metricQuery string) (types.MetricName, types.Labels, types.Aggr
 	}
 	mq := strings.TrimSpace(metricQuery)
 	aggregateFunction := types.VecSum // default
-	for _, aggFn := range []types.AggregationOverVectors{types.VecSum, types.VecAvg, types.VecMin, types.VecMax} {
+	for _, aggFn := range []types.AggregationOverVectors{types.VecSum, types.VecAvg, types.VecMin, types.VecMax, types.VecCount} {
 		if strings.HasPrefix(mq, string(aggFn)+"(") && strings.HasSuffix(mq, ")") {
 			aggregateFunction = aggFn
 			mq = strings.TrimPrefix(mq, string(aggFn)+"(")

metric/simple_parser_test.go

@@ -80,6 +80,17 @@ func TestSimpleParserMin(t *testing.T) {
 	}
 }
 
+func TestSimpleParserCount(t *testing.T) {
+	// setup
+	p := NewParser()
+
+	// check
+	name, labels, agg, err := p.Parse("count(metric_foo{ahoj=cau})")
+	if name != "metric_foo" || fmt.Sprint(labels) != fmt.Sprint(map[string]any{"ahoj": "cau"}) || agg != types.VecCount || err != nil {
+		t.Errorf("expected: [metric_foo, map[ahoj:cau], min, <nil>], got: [%s, %v, %v, %v]", name, labels, agg, err)
+	}
+}
+
 func TestSimpleParserNoLabels(t *testing.T) {
 	// setup
 	p := NewParser()

types/metrics.go

@@ -39,11 +39,16 @@ const (
 	// and on the resulting set of numbers where each represents last_one, rate, min, max, avg of the time serie, we apply
 	// this function
 
-	// VecSum sums the number
+	// VecSum sums the numbers
 	VecSum AggregationOverVectors = "sum"
+	// VecAvg calculate the mean value
 	VecAvg AggregationOverVectors = "avg"
+	// VecMin calculate the minimum value
 	VecMin AggregationOverVectors = "min"
+	// VecMax calculate the max value
 	VecMax AggregationOverVectors = "max"
+	// VecCount calculate the number of occurrences
+	VecCount AggregationOverVectors = "count"
 
 	// following operations can be applied on one time serie vector that was captured over time
 	// returning just one number