From 22e3aeb10711a69c555bfe1dee5debc2bb0de39c Mon Sep 17 00:00:00 2001
From: glutamatt <fatal974@hotmail.com>
Date: Tue, 18 Dec 2018 19:56:51 +0100
Subject: [PATCH] Add WALSegmentSize as an option of tsdb creation (#450)

Expose `WALSegmentSize` option to allow overriding the `DefaultOptions.WALSegmentSize`.
---
 CHANGELOG.md |  3 +--
 db.go        | 10 +++++++++-
 db_test.go   | 24 ++++++++++++++++++++++++
 wal/wal.go   |  4 ++--
 4 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1f2abce09..8e3c01f2c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,11 +1,10 @@
 ## master / unreleased
-
+ - [CHANGE] New `WALSegmentSize` option to override the `DefaultOptions.WALSegmentSize`. Added to allow using smaller wal files. For example using tmpfs on a RPI to minimise the SD card wear out from the constant WAL writes. As part of this change the `DefaultOptions.WALSegmentSize` constant was also exposed.
 
 ## 0.3.1
 - [BUGFIX] Fixed most windows test and some actual bugs for unclosed file readers.
 
 ## 0.3.0
-
  - [CHANGE] `LastCheckpoint()` used to return just the segment name and now it returns the full relative path.
  - [CHANGE] `NewSegmentsRangeReader()` can now read over miltiple wal ranges by using the new `SegmentRange{}` struct.
  - [CHANGE] `CorruptionErr{}` now also exposes the Segment `Dir` which is added when displaying any errors.
diff --git a/db.go b/db.go
index 3a47f0bf4..e5a057cba 100644
--- a/db.go
+++ b/db.go
@@ -45,6 +45,7 @@ import (
 // millisecond precision timestamps.
 var DefaultOptions = &Options{
 	WALFlushInterval:  5 * time.Second,
+	WALSegmentSize:    wal.DefaultSegmentSize,
 	RetentionDuration: 15 * 24 * 60 * 60 * 1000, // 15 days in milliseconds
 	BlockRanges:       ExponentialBlockRanges(int64(2*time.Hour)/1e6, 3, 5),
 	NoLockfile:        false,
@@ -55,6 +56,9 @@ type Options struct {
 	// The interval at which the write ahead log is flushed to disk.
 	WALFlushInterval time.Duration
 
+	// Segments (wal files) max size
+	WALSegmentSize int
+
 	// Duration of persisted data to keep.
 	RetentionDuration uint64
 
@@ -263,7 +267,11 @@ func Open(dir string, l log.Logger, r prometheus.Registerer, opts *Options) (db
 		return nil, errors.Wrap(err, "create leveled compactor")
 	}
 
-	wlog, err := wal.New(l, r, filepath.Join(dir, "wal"))
+	segmentSize := wal.DefaultSegmentSize
+	if opts.WALSegmentSize > 0 {
+		segmentSize = opts.WALSegmentSize
+	}
+	wlog, err := wal.NewSize(l, r, filepath.Join(dir, "wal"), segmentSize)
 	if err != nil {
 		return nil, err
 	}
diff --git a/db_test.go b/db_test.go
index 6217849d4..0ed4356cc 100644
--- a/db_test.go
+++ b/db_test.go
@@ -695,6 +695,30 @@ func TestWALFlushedOnDBClose(t *testing.T) {
 	testutil.Equals(t, []string{"labelvalue"}, values)
 }
 
+func TestWALSegmentSizeOption(t *testing.T) {
+	options := *DefaultOptions
+	options.WALSegmentSize = 2 * 32 * 1024
+	db, close := openTestDB(t, &options)
+	defer close()
+	app := db.Appender()
+	for i := int64(0); i < 155; i++ {
+		_, err := app.Add(labels.Labels{labels.Label{Name: "wal", Value: "size"}}, i, rand.Float64())
+		testutil.Ok(t, err)
+		testutil.Ok(t, app.Commit())
+	}
+
+	files, err := ioutil.ReadDir(filepath.Join(db.Dir(), "wal"))
+	testutil.Assert(t, len(files) > 1, "current WALSegmentSize should result in more than a single WAL file.")
+	testutil.Ok(t, err)
+	for i, f := range files {
+		if len(files)-1 != i {
+			testutil.Equals(t, int64(options.WALSegmentSize), f.Size(), "WAL file size doesn't match WALSegmentSize option, filename: %v", f.Name())
+			continue
+		}
+		testutil.Assert(t, int64(options.WALSegmentSize) > f.Size(), "last WAL file size is not smaller than the WALSegmentSize option, filename: %v", f.Name())
+	}
+}
+
 func TestTombstoneClean(t *testing.T) {
 	numSamples := int64(10)
 
diff --git a/wal/wal.go b/wal/wal.go
index 5134850fe..92374e312 100644
--- a/wal/wal.go
+++ b/wal/wal.go
@@ -35,7 +35,7 @@ import (
 )
 
 const (
-	defaultSegmentSize = 128 * 1024 * 1024 // 128 MB
+	DefaultSegmentSize = 128 * 1024 * 1024 // 128 MB
 	pageSize           = 32 * 1024         // 32KB
 	recordHeaderSize   = 7
 )
@@ -174,7 +174,7 @@ type WAL struct {
 
 // New returns a new WAL over the given directory.
 func New(logger log.Logger, reg prometheus.Registerer, dir string) (*WAL, error) {
-	return NewSize(logger, reg, dir, defaultSegmentSize)
+	return NewSize(logger, reg, dir, DefaultSegmentSize)
 }
 
 // NewSize returns a new WAL over the given directory.