MQE Series Metadata memory tracker

pkg/streamingpromql/operators/aggregations/aggregation.go

@@ -124,7 +124,7 @@ func (a *Aggregation) SeriesMetadata(ctx context.Context) ([]types.SeriesMetadat
 		return nil, err
 	}
 
-	defer types.PutSeriesMetadataSlice(innerSeries)
+	defer types.SeriesMetadataSlicePool.Put(innerSeries, a.MemoryConsumptionTracker)
 
 	if len(innerSeries) == 0 {
 		// No input series == no output series.
@@ -160,7 +160,11 @@ func (a *Aggregation) SeriesMetadata(ctx context.Context) ([]types.SeriesMetadat
 	}
 
 	// Sort the list of series we'll return, and maintain the order of the corresponding groups at the same time
-	seriesMetadata := types.GetSeriesMetadataSlice(len(groups))
+	seriesMetadata, err := types.SeriesMetadataSlicePool.Get(len(groups), a.MemoryConsumptionTracker)
+	if err != nil {
+		return nil, err
+	}
+
 	a.remainingGroups = make([]*group, 0, len(groups))
 
 	for _, g := range groups {

pkg/streamingpromql/operators/aggregations/count_values.go

@@ -91,7 +91,7 @@ func (c *CountValues) SeriesMetadata(ctx context.Context) ([]types.SeriesMetadat
 		return nil, err
 	}
 
-	defer types.PutSeriesMetadataSlice(innerMetadata)
+	defer types.SeriesMetadataSlicePool.Put(innerMetadata, c.MemoryConsumptionTracker)
 
 	c.labelsBuilder = labels.NewBuilder(labels.EmptyLabels())
 	c.labelsBytesBuffer = make([]byte, 0, 1024) // Why 1024 bytes? It's what labels.Labels.String() uses as a buffer size, so we use that as a sensible starting point too.
@@ -125,7 +125,11 @@ func (c *CountValues) SeriesMetadata(ctx context.Context) ([]types.SeriesMetadat
 		types.PutInstantVectorSeriesData(data, c.MemoryConsumptionTracker)
 	}
 
-	outputMetadata := types.GetSeriesMetadataSlice(len(accumulator))
+	outputMetadata, err := types.SeriesMetadataSlicePool.Get(len(accumulator), c.MemoryConsumptionTracker)
+	if err != nil {
+		return nil, err
+	}
+
 	c.series = make([][]promql.FPoint, 0, len(accumulator))
 
 	for _, s := range accumulator {

pkg/streamingpromql/operators/aggregations/topkbottomk/instant_query.go

@@ -62,7 +62,7 @@ func (t *InstantQuery) SeriesMetadata(ctx context.Context) ([]types.SeriesMetada
 		return nil, err
 	}
 
-	defer types.PutSeriesMetadataSlice(innerSeries)
+	defer types.SeriesMetadataSlicePool.Put(innerSeries, t.MemoryConsumptionTracker)
 
 	groupLabelsBytesFunc := aggregations.GroupLabelsBytesFunc(t.Grouping, t.Without)
 	groups := map[string]*instantQueryGroup{}
@@ -115,7 +115,11 @@ func (t *InstantQuery) SeriesMetadata(ctx context.Context) ([]types.SeriesMetada
 		types.PutInstantVectorSeriesData(data, t.MemoryConsumptionTracker)
 	}
 
-	outputSeries := types.GetSeriesMetadataSlice(outputSeriesCount)
+	outputSeries, err := types.SeriesMetadataSlicePool.Get(outputSeriesCount, t.MemoryConsumptionTracker)
+	if err != nil {
+		return nil, err
+	}
+
 	t.values, err = types.Float64SlicePool.Get(outputSeriesCount, t.MemoryConsumptionTracker)
 	if err != nil {
 		return nil, err

pkg/streamingpromql/operators/binops/and_unless_binary_operation.go

@@ -75,7 +75,7 @@ func (a *AndUnlessBinaryOperation) SeriesMetadata(ctx context.Context) ([]types.
 
 	if len(leftMetadata) == 0 {
 		// We can't produce any series, we are done.
-		types.PutSeriesMetadataSlice(leftMetadata)
+		types.SeriesMetadataSlicePool.Put(leftMetadata, a.MemoryConsumptionTracker)
 		return nil, nil
 	}
 
@@ -84,7 +84,7 @@ func (a *AndUnlessBinaryOperation) SeriesMetadata(ctx context.Context) ([]types.
 		return nil, err
 	}
 
-	defer types.PutSeriesMetadataSlice(rightMetadata)
+	defer types.SeriesMetadataSlicePool.Put(rightMetadata, a.MemoryConsumptionTracker)
 
 	if len(rightMetadata) == 0 && !a.IsUnless {
 		// We can't produce any series, we are done.

pkg/streamingpromql/operators/binops/grouped_vector_vector_binary_operation.go

@@ -194,7 +194,7 @@ func (g *GroupedVectorVectorBinaryOperation) SeriesMetadata(ctx context.Context)
 	}
 
 	if len(allMetadata) == 0 {
-		types.PutSeriesMetadataSlice(allMetadata)
+		types.SeriesMetadataSlicePool.Put(allMetadata, g.MemoryConsumptionTracker)
 		types.BoolSlicePool.Put(oneSideSeriesUsed, g.MemoryConsumptionTracker)
 		types.BoolSlicePool.Put(manySideSeriesUsed, g.MemoryConsumptionTracker)
 		g.Close()
@@ -383,7 +383,11 @@ func (g *GroupedVectorVectorBinaryOperation) computeOutputSeries() ([]types.Seri
 	}
 
 	// Finally, construct the list of series that this operator will return.
-	outputMetadata := types.GetSeriesMetadataSlice(len(outputSeriesMap))
+	outputMetadata, err := types.SeriesMetadataSlicePool.Get(len(outputSeriesMap), g.MemoryConsumptionTracker)
+	if err != nil {
+		return nil, nil, nil, -1, nil, -1, err
+	}
+
 	outputSeries := make([]*groupedBinaryOperationOutputSeries, 0, len(outputSeriesMap))
 
 	for _, o := range outputSeriesMap {
@@ -727,10 +731,10 @@ func (g *GroupedVectorVectorBinaryOperation) Close() {
 	g.Right.Close()
 	// We don't need to close g.oneSide or g.manySide, as these are either g.Left or g.Right and so have been closed above.
 
-	types.PutSeriesMetadataSlice(g.oneSideMetadata)
+	types.SeriesMetadataSlicePool.Put(g.oneSideMetadata, g.MemoryConsumptionTracker)
 	g.oneSideMetadata = nil
 
-	types.PutSeriesMetadataSlice(g.manySideMetadata)
+	types.SeriesMetadataSlicePool.Put(g.manySideMetadata, g.MemoryConsumptionTracker)
 	g.manySideMetadata = nil
 
 	if g.oneSideBuffer != nil {

pkg/streamingpromql/operators/binops/one_to_one_vector_vector_binary_operation.go

@@ -175,7 +175,7 @@ func (b *OneToOneVectorVectorBinaryOperation) SeriesMetadata(ctx context.Context
 	}
 
 	if len(allMetadata) == 0 {
-		types.PutSeriesMetadataSlice(allMetadata)
+		types.SeriesMetadataSlicePool.Put(allMetadata, b.MemoryConsumptionTracker)
 		types.BoolSlicePool.Put(leftSeriesUsed, b.MemoryConsumptionTracker)
 		types.BoolSlicePool.Put(rightSeriesUsed, b.MemoryConsumptionTracker)
 		b.Close()
@@ -307,7 +307,11 @@ func (b *OneToOneVectorVectorBinaryOperation) computeOutputSeries() ([]types.Ser
 		lastLeftSeriesUsedIndex = leftSeriesIndex
 	}
 
-	allMetadata := types.GetSeriesMetadataSlice(len(outputSeriesMap))
+	allMetadata, err := types.SeriesMetadataSlicePool.Get(len(outputSeriesMap), b.MemoryConsumptionTracker)
+	if err != nil {
+		return nil, nil, nil, -1, nil, -1, err
+	}
+
 	allSeries := make([]*oneToOneBinaryOperationOutputSeries, 0, len(outputSeriesMap))
 
 	for _, outputSeries := range outputSeriesMap {
@@ -568,10 +572,10 @@ func (b *OneToOneVectorVectorBinaryOperation) Close() {
 	b.Left.Close()
 	b.Right.Close()
 
-	types.PutSeriesMetadataSlice(b.leftMetadata)
+	types.SeriesMetadataSlicePool.Put(b.leftMetadata, b.MemoryConsumptionTracker)
 	b.leftMetadata = nil
 
-	types.PutSeriesMetadataSlice(b.rightMetadata)
+	types.SeriesMetadataSlicePool.Put(b.rightMetadata, b.MemoryConsumptionTracker)
 	b.rightMetadata = nil
 
 	if b.leftBuffer != nil {

pkg/streamingpromql/operators/binops/or_binary_operation.go

@@ -73,8 +73,8 @@ func (o *OrBinaryOperation) SeriesMetadata(ctx context.Context) ([]types.SeriesM
 
 	if len(leftMetadata) == 0 && len(rightMetadata) == 0 {
 		// Nothing to return.
-		types.PutSeriesMetadataSlice(leftMetadata)
-		types.PutSeriesMetadataSlice(rightMetadata)
+		types.SeriesMetadataSlicePool.Put(leftMetadata, o.MemoryConsumptionTracker)
+		types.SeriesMetadataSlicePool.Put(rightMetadata, o.MemoryConsumptionTracker)
 
 		o.Left.Close()
 		o.Right.Close()
@@ -86,7 +86,7 @@ func (o *OrBinaryOperation) SeriesMetadata(ctx context.Context) ([]types.SeriesM
 		// We can just return everything from the right side.
 		o.nextSeriesIsFromLeft = false
 		o.rightSeriesCount = []int{len(rightMetadata)}
-		types.PutSeriesMetadataSlice(leftMetadata)
+		types.SeriesMetadataSlicePool.Put(leftMetadata, o.MemoryConsumptionTracker)
 
 		o.Left.Close()
 
@@ -97,19 +97,19 @@ func (o *OrBinaryOperation) SeriesMetadata(ctx context.Context) ([]types.SeriesM
 		// We can just return everything from the left side.
 		o.nextSeriesIsFromLeft = true
 		o.leftSeriesCount = []int{len(leftMetadata)}
-		types.PutSeriesMetadataSlice(rightMetadata)
+		types.SeriesMetadataSlicePool.Put(rightMetadata, o.MemoryConsumptionTracker)
 
 		o.Right.Close()
 
 		return leftMetadata, nil
 	}
 
-	defer types.PutSeriesMetadataSlice(leftMetadata)
-	defer types.PutSeriesMetadataSlice(rightMetadata)
+	defer types.SeriesMetadataSlicePool.Put(leftMetadata, o.MemoryConsumptionTracker)
+	defer types.SeriesMetadataSlicePool.Put(rightMetadata, o.MemoryConsumptionTracker)
 
 	o.computeGroups(leftMetadata, rightMetadata)
 
-	return o.computeSeriesOutputOrder(leftMetadata, rightMetadata), nil
+	return o.computeSeriesOutputOrder(leftMetadata, rightMetadata)
 }
 
 func (o *OrBinaryOperation) computeGroups(leftMetadata []types.SeriesMetadata, rightMetadata []types.SeriesMetadata) {
@@ -155,7 +155,7 @@ func (o *OrBinaryOperation) computeGroups(leftMetadata []types.SeriesMetadata, r
 	}
 }
 
-func (o *OrBinaryOperation) computeSeriesOutputOrder(leftMetadata []types.SeriesMetadata, rightMetadata []types.SeriesMetadata) []types.SeriesMetadata {
+func (o *OrBinaryOperation) computeSeriesOutputOrder(leftMetadata []types.SeriesMetadata, rightMetadata []types.SeriesMetadata) ([]types.SeriesMetadata, error) {
 	// The idea here is to determine the order we should return series in, returning series from the right side as soon as we've seen all
 	// the series from the left that we need.
 	//
@@ -174,7 +174,10 @@ func (o *OrBinaryOperation) computeSeriesOutputOrder(leftMetadata []types.Series
 	// state on both sides.
 
 	nextLeftSeriesToRead := 0
-	series := types.GetSeriesMetadataSlice(len(leftMetadata) + len(rightMetadata))
+	series, err := types.SeriesMetadataSlicePool.Get(len(leftMetadata)+len(rightMetadata), o.MemoryConsumptionTracker)
+	if err != nil {
+		return nil, err
+	}
 
 	for nextRightSeriesToRead, rightGroup := range o.rightSeriesGroups {
 		lastSeriesFromLeft := false
@@ -215,7 +218,7 @@ func (o *OrBinaryOperation) computeSeriesOutputOrder(leftMetadata []types.Series
 		o.leftSeriesCount = append(o.leftSeriesCount, seriesCount)
 	}
 
-	return series
+	return series, nil
 }
 
 func (o *OrBinaryOperation) NextSeries(ctx context.Context) (types.InstantVectorSeriesData, error) {

pkg/streamingpromql/operators/deduplicate_and_merge.go

@@ -50,16 +50,20 @@ func (d *DeduplicateAndMerge) SeriesMetadata(ctx context.Context) ([]types.Serie
 	}
 
 	// We might have duplicates (or HasDuplicateSeries hit a hash collision). Determine the merged output series.
-	groups, outputMetadata := d.computeOutputSeriesGroups(innerMetadata)
+	groups, outputMetadata, err := d.computeOutputSeriesGroups(innerMetadata)
+	if err != nil {
+		return nil, err
+	}
+
 	d.groups = groups
-	types.PutSeriesMetadataSlice(innerMetadata)
+	types.SeriesMetadataSlicePool.Put(innerMetadata, d.MemoryConsumptionTracker)
 
 	d.buffer = NewInstantVectorOperatorBuffer(d.Inner, nil, len(innerMetadata), d.MemoryConsumptionTracker)
 
 	return outputMetadata, nil
 }
 
-func (d *DeduplicateAndMerge) computeOutputSeriesGroups(innerMetadata []types.SeriesMetadata) ([][]int, []types.SeriesMetadata) {
+func (d *DeduplicateAndMerge) computeOutputSeriesGroups(innerMetadata []types.SeriesMetadata) ([][]int, []types.SeriesMetadata, error) {
 	// Why use a string, rather than the labels hash as a key here? This avoids any issues with hash collisions.
 	outputGroupMap := map[string][]int{}
 
@@ -91,13 +95,16 @@ func (d *DeduplicateAndMerge) computeOutputSeriesGroups(innerMetadata []types.Se
 	})
 
 	// Now that we know which series we'll return, and in what order, create the list of output series.
-	outputMetadata := types.GetSeriesMetadataSlice(len(outputGroups))
+	outputMetadata, err := types.SeriesMetadataSlicePool.Get(len(outputGroups), d.MemoryConsumptionTracker)
+	if err != nil {
+		return nil, nil, err
+	}
 
 	for _, group := range outputGroups {
 		outputMetadata = append(outputMetadata, innerMetadata[group[0]])
 	}
 
-	return outputGroups, outputMetadata
+	return outputGroups, outputMetadata, nil
 }
 
 func (d *DeduplicateAndMerge) NextSeries(ctx context.Context) (types.InstantVectorSeriesData, error) {

pkg/streamingpromql/operators/functions/absent.go

@@ -47,7 +47,7 @@ func (a *Absent) SeriesMetadata(ctx context.Context) ([]types.SeriesMetadata, er
 	if err != nil {
 		return nil, err
 	}
-	defer types.PutSeriesMetadataSlice(innerMetadata)
+	defer types.SeriesMetadataSlicePool.Put(innerMetadata, a.MemoryConsumptionTracker)
 
 	a.presence, err = types.BoolSlicePool.Get(a.TimeRange.StepCount, a.MemoryConsumptionTracker)
 	if err != nil {
@@ -57,7 +57,11 @@ func (a *Absent) SeriesMetadata(ctx context.Context) ([]types.SeriesMetadata, er
 	// Initialize presence slice
 	a.presence = a.presence[:a.TimeRange.StepCount]
 
-	metadata := types.GetSeriesMetadataSlice(1)
+	metadata, err := types.SeriesMetadataSlicePool.Get(1, a.MemoryConsumptionTracker)
+	if err != nil {
+		return nil, err
+	}
+
 	metadata = append(metadata, types.SeriesMetadata{
 		Labels: a.Labels,
 	})

pkg/streamingpromql/operators/functions/absent_over_time.go

@@ -51,7 +51,7 @@ func (a *AbsentOverTime) SeriesMetadata(ctx context.Context) ([]types.SeriesMeta
 	if err != nil {
 		return nil, err
 	}
-	defer types.PutSeriesMetadataSlice(innerMetadata)
+	defer types.SeriesMetadataSlicePool.Put(innerMetadata, a.MemoryConsumptionTracker)
 
 	a.presence, err = types.BoolSlicePool.Get(a.TimeRange.StepCount, a.MemoryConsumptionTracker)
 	if err != nil {
@@ -61,7 +61,11 @@ func (a *AbsentOverTime) SeriesMetadata(ctx context.Context) ([]types.SeriesMeta
 	// Initialize presence slice
 	a.presence = a.presence[:a.TimeRange.StepCount]
 
-	metadata := types.GetSeriesMetadataSlice(1)
+	metadata, err := types.SeriesMetadataSlicePool.Get(1, a.MemoryConsumptionTracker)
+	if err != nil {
+		return nil, err
+	}
+
 	metadata = append(metadata, types.SeriesMetadata{
 		Labels: a.Labels,
 	})

pkg/streamingpromql/operators/functions/factories.go

@@ -86,7 +86,7 @@ func TimeTransformationFunctionOperatorFactory(name string, seriesDataFunc Insta
 		var inner types.InstantVectorOperator
 		if len(args) == 0 {
 			// if the argument is not provided, it will default to vector(time())
-			inner = scalars.NewScalarToInstantVector(operators.NewTime(timeRange, memoryConsumptionTracker, expressionPosition), expressionPosition)
+			inner = scalars.NewScalarToInstantVector(operators.NewTime(timeRange, memoryConsumptionTracker, expressionPosition), expressionPosition, memoryConsumptionTracker)
 		} else if len(args) == 1 {
 			// if one argument is provided, it must be an instant vector
 			var ok bool
@@ -216,7 +216,7 @@ func QuantileOverTimeFactory(args []types.Operator, memoryConsumptionTracker *li
 	return o, nil
 }
 
-func scalarToInstantVectorOperatorFactory(args []types.Operator, _ *limiting.MemoryConsumptionTracker, _ *annotations.Annotations, expressionPosition posrange.PositionRange, _ types.QueryTimeRange) (types.InstantVectorOperator, error) {
+func scalarToInstantVectorOperatorFactory(args []types.Operator, memoryConsumptionTracker *limiting.MemoryConsumptionTracker, _ *annotations.Annotations, expressionPosition posrange.PositionRange, _ types.QueryTimeRange) (types.InstantVectorOperator, error) {
 	if len(args) != 1 {
 		// Should be caught by the PromQL parser, but we check here for safety.
 		return nil, fmt.Errorf("expected exactly 1 argument for vector, got %v", len(args))
@@ -228,7 +228,7 @@ func scalarToInstantVectorOperatorFactory(args []types.Operator, _ *limiting.Mem
 		return nil, fmt.Errorf("expected a scalar argument for vector, got %T", args[0])
 	}
 
-	return scalars.NewScalarToInstantVector(inner, expressionPosition), nil
+	return scalars.NewScalarToInstantVector(inner, expressionPosition, memoryConsumptionTracker), nil
 }
 
 func LabelJoinFunctionOperatorFactory(args []types.Operator, memoryConsumptionTracker *limiting.MemoryConsumptionTracker, _ *annotations.Annotations, expressionPosition posrange.PositionRange, timeRange types.QueryTimeRange) (types.InstantVectorOperator, error) {

pkg/streamingpromql/operators/functions/histogram_quantile.go

@@ -141,7 +141,7 @@ func (h *HistogramQuantileFunction) SeriesMetadata(ctx context.Context) ([]types
 	if err != nil {
 		return nil, err
 	}
-	defer types.PutSeriesMetadataSlice(innerSeries)
+	defer types.SeriesMetadataSlicePool.Put(innerSeries, h.memoryConsumptionTracker)
 
 	if len(innerSeries) == 0 {
 		// No input series == no output series.
@@ -197,7 +197,11 @@ func (h *HistogramQuantileFunction) SeriesMetadata(ctx context.Context) ([]types
 		h.seriesGroupPairs[innerIdx].classicHistogramGroup = g.group
 	}
 
-	seriesMetadata := types.GetSeriesMetadataSlice(len(groups))
+	seriesMetadata, err := types.SeriesMetadataSlicePool.Get(len(groups), h.memoryConsumptionTracker)
+	if err != nil {
+		return nil, err
+	}
+
 	h.remainingGroups = make([]*bucketGroup, 0, len(groups))
 	for _, g := range groups {
 		seriesMetadata = append(seriesMetadata, types.SeriesMetadata{Labels: g.labels.DropMetricName()})

pkg/streamingpromql/operators/scalars/instant_vector_to_scalar.go

@@ -96,7 +96,7 @@ func (i *InstantVectorToScalar) getInnerSeriesCount(ctx context.Context) (int, e
 		return 0, err
 	}
 
-	defer types.PutSeriesMetadataSlice(metadata)
+	defer types.SeriesMetadataSlicePool.Put(metadata, i.MemoryConsumptionTracker)
 
 	seriesCount := len(metadata)