From 0fa3a72e84b5b858af63c2fafe7741de1cdc4fd3 Mon Sep 17 00:00:00 2001
From: machine424 <ayoubmrini424@gmail.com>
Date: Mon, 30 Dec 2024 21:35:09 +0100
Subject: [PATCH] feat(tsdb/chunkenc): use pools for iterators to help reduce
 allocations under heavy usage of Appender()

The initialization of a new iterator is no longer systematic.

This is part of the "reduce WAL replay overhead/garbage" effort to help
with https://github.com/prometheus/prometheus/issues/6934.

A xorIterator allocates ~100 bytes on average per head chunk ~ per time series.
histogramIterator and floatHistogramIterator have even more overhead

During WAL replay with 10M series, `xorIterator` could generate more than 1GB of garbage.

Signed-off-by: machine424 <ayoubmrini424@gmail.com>
---
 tsdb/chunkenc/float_histogram.go | 12 +++++++++++-
 tsdb/chunkenc/histogram.go       | 12 +++++++++++-
 tsdb/chunkenc/xor.go             | 11 ++++++++++-
 3 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/tsdb/chunkenc/float_histogram.go b/tsdb/chunkenc/float_histogram.go
index e5ad4028bb..301eb55ac6 100644
--- a/tsdb/chunkenc/float_histogram.go
+++ b/tsdb/chunkenc/float_histogram.go
@@ -18,11 +18,20 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"sync"
 
 	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/value"
 )
 
+// floatHistogramChunkAppenderIterPool is exclusively used in Appender() to help reduce
+// allocations under heavy usage.
+var floatHistogramChunkAppenderIterPool = sync.Pool{
+	New: func() interface{} {
+		return &floatHistogramIterator{}
+	},
+}
+
 // FloatHistogramChunk holds encoded sample data for a sparse, high-resolution
 // float histogram.
 //
@@ -104,7 +113,8 @@ func (c *FloatHistogramChunk) Compact() {
 
 // Appender implements the Chunk interface.
 func (c *FloatHistogramChunk) Appender() (Appender, error) {
-	it := c.iterator(nil)
+	it := c.iterator(floatHistogramChunkAppenderIterPool.Get().(*floatHistogramIterator))
+	defer floatHistogramChunkAppenderIterPool.Put(it)
 
 	// To get an appender, we must know the state it would have if we had
 	// appended all existing data from scratch. We iterate through the end
diff --git a/tsdb/chunkenc/histogram.go b/tsdb/chunkenc/histogram.go
index 0f54eb6928..2b4390f099 100644
--- a/tsdb/chunkenc/histogram.go
+++ b/tsdb/chunkenc/histogram.go
@@ -18,11 +18,20 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"sync"
 
 	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/value"
 )
 
+// histogramChunkAppenderIterPool is exclusively used in Appender() to help reduce
+// allocations under heavy usage.
+var histogramChunkAppenderIterPool = sync.Pool{
+	New: func() interface{} {
+		return &histogramIterator{}
+	},
+}
+
 // HistogramChunk holds encoded sample data for a sparse, high-resolution
 // histogram.
 //
@@ -115,7 +124,8 @@ func (c *HistogramChunk) Compact() {
 
 // Appender implements the Chunk interface.
 func (c *HistogramChunk) Appender() (Appender, error) {
-	it := c.iterator(nil)
+	it := c.iterator(histogramChunkAppenderIterPool.Get().(*histogramIterator))
+	defer histogramChunkAppenderIterPool.Put(it)
 
 	// To get an appender, we must know the state it would have if we had
 	// appended all existing data from scratch. We iterate through the end
diff --git a/tsdb/chunkenc/xor.go b/tsdb/chunkenc/xor.go
index ac75a5994b..db0ffd3a6d 100644
--- a/tsdb/chunkenc/xor.go
+++ b/tsdb/chunkenc/xor.go
@@ -47,6 +47,7 @@ import (
 	"encoding/binary"
 	"math"
 	"math/bits"
+	"sync"
 
 	"github.com/prometheus/prometheus/model/histogram"
 )
@@ -55,6 +56,13 @@ const (
 	chunkCompactCapacityThreshold = 32
 )
 
+// xorAppenderIterPool is exclusively used in Appender() to help reduce allocations under heavy usage.
+var xorAppenderIterPool = sync.Pool{
+	New: func() interface{} {
+		return &xorIterator{}
+	},
+}
+
 // XORChunk holds XOR encoded sample data.
 type XORChunk struct {
 	b bstream
@@ -98,7 +106,8 @@ func (c *XORChunk) Compact() {
 // It is not valid to call Appender() multiple times concurrently or to use multiple
 // Appenders on the same chunk.
 func (c *XORChunk) Appender() (Appender, error) {
-	it := c.iterator(nil)
+	it := c.iterator(xorAppenderIterPool.Get().(*xorIterator))
+	defer xorAppenderIterPool.Put(it)
 
 	// To get an appender we must know the state it would have if we had
 	// appended all existing data from scratch.