diff --git a/tsdb/chunks/chunks.go b/tsdb/chunks/chunks.go index d89efe31b7..034106238e 100644 --- a/tsdb/chunks/chunks.go +++ b/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) diff --git a/tsdb/chunks/chunks_test.go b/tsdb/chunks/chunks_test.go index 86f2fbddbd..6eb00f12ad 100644 --- a/tsdb/chunks/chunks_test.go +++ b/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") +}