Allow human-readable byte sizes in config

Update support in the `toml` package for parsing human-readble byte sizes.
Supported size suffixes are "k" or "K" for kibibytes, "m" or "M" for
mebibytes, and "g" or "G" for gibibytes. If a size suffix isn't specified
then bytes are assumed.

In the config, `cache-max-memory-size` and `cache-snapshot-memory-size` are
now typed as `toml.Size` and support the new syntax.

CHANGELOG.md

@@ -46,6 +46,7 @@
 - [#8893](https://github.com/influxdata/influxdb/pull/8893): Handle nil MeasurementIterator.
 - [#8986](https://github.com/influxdata/influxdb/issues/8986): Add long-line support to client importer. Thanks @lets00!
 - [#9021](https://github.com/influxdata/influxdb/pull/9021): Update to go 1.9.2
+- [#8891](https://github.com/influxdata/influxdb/pull/8891): Allow human-readable byte sizes in config
 
 ### Bugfixes
 

etc/config.sample.toml

@@ -71,10 +71,12 @@
 
   # CacheMaxMemorySize is the maximum size a shard's cache can
   # reach before it starts rejecting writes.
+  # Values without a size suffix (e.g. 30k, 20m, 10g) are in bytes.
   # cache-max-memory-size = 1048576000
 
   # CacheSnapshotMemorySize is the size at which the engine will
   # snapshot the cache and write it to a TSM file, freeing up memory
+  # Values without a size suffix (e.g. 30k, 20m, 10g) are in bytes.
   # cache-snapshot-memory-size = 26214400
 
   # CacheSnapshotWriteColdDuration is the length of time at

toml/toml.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"strconv"
 	"time"
+	"unicode"
 )
 
 // maxInt is the largest integer representable by a word (architecture dependent).
@@ -42,31 +43,54 @@ func (d Duration) MarshalText() (text []byte, err error) {
 }
 
 // Size represents a TOML parseable file size.
-// Users can specify size using "m" for megabytes and "g" for gigabytes.
-type Size int
+// Users can specify size using "k" or "K" for kibibytes, "m" or "M" for mebibytes,
+// and "g" or "G" for gibibytes. If a size suffix isn't specified then bytes are assumed.
+type Size uint64
 
 // UnmarshalText parses a byte size from text.
 func (s *Size) UnmarshalText(text []byte) error {
-	// Parse numeric portion of value.
-	length := len(string(text))
-	size, err := strconv.ParseInt(string(text[:length-1]), 10, 64)
-	if err != nil {
-		return err
+	if len(text) == 0 {
+		return fmt.Errorf("size was empty")
 	}
 
-	// Parse unit of measure ("m", "g", etc).
-	switch suffix := text[len(text)-1]; suffix {
-	case 'm':
-		size *= 1 << 20 // MB
-	case 'g':
-		size *= 1 << 30 // GB
-	default:
-		return fmt.Errorf("unknown size suffix: %c", suffix)
+	// The multiplier defaults to 1 in case the size has
+	// no suffix (and is then just raw bytes)
+	mult := int64(1)
+
+	// Preserve the original text for error messages
+	sizeText := text
+
+	// Parse unit of measure
+	suffix := text[len(sizeText)-1]
+	if !unicode.IsDigit(rune(suffix)) {
+		switch suffix {
+		case 'k', 'K':
+			mult = 1 << 10 // KiB
+		case 'm', 'M':
+			mult = 1 << 20 // MiB
+		case 'g', 'G':
+			mult = 1 << 30 // GiB
+		default:
+			return fmt.Errorf("unknown size suffix: %c", suffix)
+		}
+		sizeText = sizeText[:len(sizeText)-1]
 	}
 
+	// Parse numeric portion of value.
+	size, err := strconv.ParseInt(string(sizeText), 10, 64)
+	if err != nil {
+		return fmt.Errorf("invalid size: %s", string(text))
+	}
+
+	if maxInt/mult < size {
+		return fmt.Errorf("size would overflow the max size (%d) of an int: %s", maxInt, string(text))
+	}
+
+	size *= mult
+
 	// Check for overflow.
 	if size > maxInt {
-		return fmt.Errorf("size %d cannot be represented by an int", size)
+		return fmt.Errorf("size %d is too large", size)
 	}
 
 	*s = Size(size)

toml/toml_test.go

@@ -2,6 +2,7 @@ package toml_test
 
 import (
 	"bytes"
+	"fmt"
 	"strings"
 	"testing"
 	"time"
@@ -11,23 +12,54 @@ import (
 	itoml "github.com/influxdata/influxdb/toml"
 )
 
-// Ensure that megabyte sizes can be parsed.
-func TestSize_UnmarshalText_MB(t *testing.T) {
-	var s itoml.Size
-	if err := s.UnmarshalText([]byte("200m")); err != nil {
-		t.Fatalf("unexpected error: %s", err)
-	} else if s != 200*(1<<20) {
-		t.Fatalf("unexpected size: %d", s)
-	}
-}
+func TestSize_UnmarshalText(t *testing.T) {
+	// Copy this from the toml package
+	maxInt := int64(^uint(0) >> 1)
 
-// Ensure that gigabyte sizes can be parsed.
-func TestSize_UnmarshalText_GB(t *testing.T) {
 	var s itoml.Size
-	if err := s.UnmarshalText([]byte("1g")); err != nil {
-		t.Fatalf("unexpected error: %s", err)
-	} else if s != 1073741824 {
-		t.Fatalf("unexpected size: %d", s)
+	for _, test := range []struct {
+		str  string
+		want int64
+	}{
+		{"1", 1},
+		{"10", 10},
+		{"100", 100},
+		{"1k", 1 << 10},
+		{"10k", 10 << 10},
+		{"100k", 100 << 10},
+		{"1K", 1 << 10},
+		{"10K", 10 << 10},
+		{"100K", 100 << 10},
+		{"1m", 1 << 20},
+		{"10m", 10 << 20},
+		{"100m", 100 << 20},
+		{"1M", 1 << 20},
+		{"10M", 10 << 20},
+		{"100M", 100 << 20},
+		{"1g", 1 << 30},
+		{"1G", 1 << 30},
+		{fmt.Sprint(maxInt - 1), maxInt - 1},
+	} {
+		if err := s.UnmarshalText([]byte(test.str)); err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+		if s != itoml.Size(test.want) {
+			t.Fatalf("wanted: %d got: %d", test.want, s)
+		}
+	}
+
+	for _, str := range []string{
+		fmt.Sprintf("%dk", maxInt-1),
+		"10000000000000000000g",
+		"abcdef",
+		"1KB",
+		"√m",
+		"a1",
+		"",
+	} {
+		if err := s.UnmarshalText([]byte(str)); err == nil {
+			t.Fatalf("input should have failed: %s", str)
+		}
 	}
 }
 

tsdb/config.go

@@ -71,8 +71,8 @@ type Config struct {
 	QueryLogEnabled bool `toml:"query-log-enabled"`
 
 	// Compaction options for tsm1 (descriptions above with defaults)
-	CacheMaxMemorySize             uint64        `toml:"cache-max-memory-size"`
-	CacheSnapshotMemorySize        uint64        `toml:"cache-snapshot-memory-size"`
+	CacheMaxMemorySize             toml.Size     `toml:"cache-max-memory-size"`
+	CacheSnapshotMemorySize        toml.Size     `toml:"cache-snapshot-memory-size"`
 	CacheSnapshotWriteColdDuration toml.Duration `toml:"cache-snapshot-write-cold-duration"`
 	CompactFullWriteColdDuration   toml.Duration `toml:"compact-full-write-cold-duration"`
 
@@ -105,8 +105,8 @@ func NewConfig() Config {
 
 		QueryLogEnabled: true,
 
-		CacheMaxMemorySize:             DefaultCacheMaxMemorySize,
-		CacheSnapshotMemorySize:        DefaultCacheSnapshotMemorySize,
+		CacheMaxMemorySize:             toml.Size(DefaultCacheMaxMemorySize),
+		CacheSnapshotMemorySize:        toml.Size(DefaultCacheSnapshotMemorySize),
 		CacheSnapshotWriteColdDuration: toml.Duration(DefaultCacheSnapshotWriteColdDuration),
 		CompactFullWriteColdDuration:   toml.Duration(DefaultCompactFullWriteColdDuration),
 

tsdb/engine/tsm1/engine.go

@@ -194,7 +194,7 @@ func NewEngine(id uint64, idx tsdb.Index, database, path string, walPath string,
 		Compactor:      c,
 		CompactionPlan: NewDefaultPlanner(fs, time.Duration(opt.Config.CompactFullWriteColdDuration)),
 
-		CacheFlushMemorySizeThreshold: opt.Config.CacheSnapshotMemorySize,
+		CacheFlushMemorySizeThreshold: uint64(opt.Config.CacheSnapshotMemorySize),
 		CacheFlushWriteColdDuration:   time.Duration(opt.Config.CacheSnapshotWriteColdDuration),
 		enableCompactionsOnOpen:       true,
 		stats:             stats,