From 4696b46dd5ef89e4e23554dfec9e838c6bbaf23e Mon Sep 17 00:00:00 2001
From: beorn7 <beorn@grafana.com>
Date: Tue, 31 Oct 2023 14:50:26 +0100
Subject: [PATCH] storage: Fix mixed samples handling in sampleRing

Two issues are fixed here, that lead to the same problem:

1. If `newSampleRing` is called with an unknown ValueType including
   ValueNone, we have initialized the interface buffer (`iBuf`).
   However, we would still use a specialized buffer for the first
   sample, opportunistically assuming that we might still not
   encounter mixed samples and we should go down the more efficient
   road.

2. If the `sampleRing` is `reset`, we leave all buffers alone,
   including `iBuf`, which is generally fine, but not for `iBuf`, see
   below.

In both cases, `iBuf` already contains values, but we will fill one of
the specialized buffers first. Once we then actually encounter mixed
samples, the content of the specialized buffer is copied into `iBuf`
using `append`. That's by itself the right idea because `iBuf` might
be `nil`, and even if not, it might or might not have the right
capacity. However, this approach assumes that `iBuf` is empty, or more
precisely has a length of zero.

This commit makes sure that `iBuf` does not get needlessly initialized
in `newSampleRing` and that it is emptied upon `reset`.

A test case is added to demonstrate both issues above.

Signed-off-by: beorn7 <beorn@grafana.com>
---
 storage/buffer.go      | 10 ++++++++-
 storage/buffer_test.go | 48 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/storage/buffer.go b/storage/buffer.go
index a692570ed..d2d89e042 100644
--- a/storage/buffer.go
+++ b/storage/buffer.go
@@ -284,7 +284,8 @@ func newSampleRing(delta int64, size int, typ chunkenc.ValueType) *sampleRing {
 	case chunkenc.ValFloatHistogram:
 		r.fhBuf = make([]fhSample, size)
 	default:
-		r.iBuf = make([]chunks.Sample, size)
+		// Do not initialize anything because the 1st sample will be
+		// added to one of the other bufs anyway.
 	}
 	return r
 }
@@ -294,6 +295,12 @@ func (r *sampleRing) reset() {
 	r.i = -1
 	r.f = 0
 	r.bufInUse = noBuf
+
+	// The first sample after the reset will always go to a specialized
+	// buffer. If we later need to change to the interface buffer, we'll
+	// copy from the specialized buffer to the interface buffer. For that to
+	// work properly, we have to reset the interface buffer here, too.
+	r.iBuf = r.iBuf[:0]
 }
 
 // Returns the current iterator. Invalidates previously returned iterators.
@@ -441,6 +448,7 @@ func (r *sampleRing) add(s chunks.Sample) {
 		}
 		// The new sample isn't a fit for the already existing
 		// ones. Copy the latter into the interface buffer where needed.
+		// The interface buffer is assumed to be of length zero at this point.
 		switch r.bufInUse {
 		case fBuf:
 			for _, s := range r.fBuf {
diff --git a/storage/buffer_test.go b/storage/buffer_test.go
index bc79a79ba..77981545b 100644
--- a/storage/buffer_test.go
+++ b/storage/buffer_test.go
@@ -90,6 +90,54 @@ func TestSampleRing(t *testing.T) {
 	}
 }
 
+func TestSampleRingMixed(t *testing.T) {
+	h1 := tsdbutil.GenerateTestHistogram(1)
+	h2 := tsdbutil.GenerateTestHistogram(2)
+
+	// With ValNone as the preferred type, nothing should be initialized.
+	r := newSampleRing(10, 2, chunkenc.ValNone)
+	require.Zero(t, len(r.fBuf))
+	require.Zero(t, len(r.hBuf))
+	require.Zero(t, len(r.fhBuf))
+	require.Zero(t, len(r.iBuf))
+
+	// But then mixed adds should work as expected.
+	r.addF(fSample{t: 1, f: 3.14})
+	r.addH(hSample{t: 2, h: h1})
+
+	it := r.iterator()
+
+	require.Equal(t, chunkenc.ValFloat, it.Next())
+	ts, f := it.At()
+	require.Equal(t, int64(1), ts)
+	require.Equal(t, 3.14, f)
+	require.Equal(t, chunkenc.ValHistogram, it.Next())
+	var h *histogram.Histogram
+	ts, h = it.AtHistogram()
+	require.Equal(t, int64(2), ts)
+	require.Equal(t, h1, h)
+	require.Equal(t, chunkenc.ValNone, it.Next())
+
+	r.reset()
+	it = r.iterator()
+	require.Equal(t, chunkenc.ValNone, it.Next())
+
+	r.addF(fSample{t: 3, f: 4.2})
+	r.addH(hSample{t: 4, h: h2})
+
+	it = r.iterator()
+
+	require.Equal(t, chunkenc.ValFloat, it.Next())
+	ts, f = it.At()
+	require.Equal(t, int64(3), ts)
+	require.Equal(t, 4.2, f)
+	require.Equal(t, chunkenc.ValHistogram, it.Next())
+	ts, h = it.AtHistogram()
+	require.Equal(t, int64(4), ts)
+	require.Equal(t, h2, h)
+	require.Equal(t, chunkenc.ValNone, it.Next())
+}
+
 func TestBufferedSeriesIterator(t *testing.T) {
 	var it *BufferedSeriesIterator