tsdb: fix issue where a new segment file is created for every chunk if WithSegmentSize not called (#16635)

* tsdb: fix issue where a new segment file is created for every chunk

Signed-off-by: Charles Korn <charles.korn@grafana.com>

* Address PR feedback

Signed-off-by: Charles Korn <charles.korn@grafana.com>

---------

Signed-off-by: Charles Korn <charles.korn@grafana.com>

tsdb/chunks/chunks.go

@@ -307,18 +307,24 @@ func WithUncachedIO(enabled bool) WriterOption {
 	}
 }
 
+// WithSegmentSize sets the chunk segment size for the writer.
+// Passing a value less than or equal to 0 causes the default segment size (DefaultChunkSegmentSize) to be used.
 func WithSegmentSize(segmentSize int64) WriterOption {
 	return func(o *writerOptions) {
 		if segmentSize <= 0 {
 			segmentSize = DefaultChunkSegmentSize
 		}
+
 		o.segmentSize = segmentSize
 	}
 }
 
 // NewWriter returns a new writer against the given directory.
+// It uses DefaultChunkSegmentSize as the default segment size.
 func NewWriter(dir string, opts ...WriterOption) (*Writer, error) {
-	options := &writerOptions{}
+	options := &writerOptions{
+		segmentSize: DefaultChunkSegmentSize,
+	}
 
 	for _, opt := range opts {
 		opt(options)

tsdb/chunks/chunks_test.go

@@ -14,9 +14,12 @@
 package chunks
 
 import (
+	"os"
 	"testing"
 
 	"github.com/stretchr/testify/require"
+
+	"github.com/prometheus/prometheus/tsdb/tsdbutil"
 )
 
 func TestReaderWithInvalidBuffer(t *testing.T) {
@@ -26,3 +29,32 @@ func TestReaderWithInvalidBuffer(t *testing.T) {
 	_, _, err := r.ChunkOrIterable(Meta{Ref: 0})
 	require.Error(t, err)
 }
+
+func TestWriterWithDefaultSegmentSize(t *testing.T) {
+	chk1, err := ChunkFromSamples([]Sample{
+		sample{t: 10, f: 11},
+		sample{t: 20, f: 12},
+		sample{t: 30, f: 13},
+	})
+	require.NoError(t, err)
+
+	chk2, err := ChunkFromSamples([]Sample{
+		sample{t: 40, h: tsdbutil.GenerateTestHistogram(1)},
+		sample{t: 50, h: tsdbutil.GenerateTestHistogram(2)},
+		sample{t: 60, h: tsdbutil.GenerateTestHistogram(3)},
+	})
+	require.NoError(t, err)
+
+	dir := t.TempDir()
+	w, err := NewWriter(dir)
+	require.NoError(t, err)
+
+	err = w.WriteChunks(chk1, chk2)
+	require.NoError(t, err)
+
+	require.NoError(t, w.Close())
+
+	d, err := os.ReadDir(dir)
+	require.NoError(t, err)
+	require.Len(t, d, 1, "expected only one segment to be created to hold both chunks")
+}