Add ability to specify bytes options as strings with units (KB, MiB, ...) (#4852)

internal/internal.go

@@ -16,6 +16,8 @@ import (
 	"syscall"
 	"time"
 	"unicode"
+
+	"github.com/alecthomas/units"
 )
 
 const alphanum string = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
@@ -36,6 +38,11 @@ type Duration struct {
 	Duration time.Duration
 }
 
+// Size just wraps an int64
+type Size struct {
+	Size int64
+}
+
 // SetVersion sets the telegraf agent version
 func SetVersion(v string) error {
 	if version != "" {
@@ -85,6 +92,27 @@ func (d *Duration) UnmarshalTOML(b []byte) error {
 	return nil
 }
 
+func (s *Size) UnmarshalTOML(b []byte) error {
+	var err error
+	b = bytes.Trim(b, `'`)
+
+	val, err := strconv.ParseInt(string(b), 10, 64)
+	if err == nil {
+		s.Size = val
+		return nil
+	}
+	uq, err := strconv.Unquote(string(b))
+	if err != nil {
+		return err
+	}
+	val, err = units.ParseStrictBytes(uq)
+	if err != nil {
+		return err
+	}
+	s.Size = val
+	return nil
+}
+
 // ReadLines reads contents from a file and splits them by new lines.
 // A convenience wrapper to ReadLinesOffsetN(filename, 0, -1).
 func ReadLines(filename string) ([]string, error) {

internal/internal_test.go

@@ -166,6 +166,29 @@ func TestDuration(t *testing.T) {
 	assert.Equal(t, time.Second, d.Duration)
 }
 
+func TestSize(t *testing.T) {
+	var s Size
+
+	s.UnmarshalTOML([]byte(`"1B"`))
+	assert.Equal(t, int64(1), s.Size)
+
+	s = Size{}
+	s.UnmarshalTOML([]byte(`1`))
+	assert.Equal(t, int64(1), s.Size)
+
+	s = Size{}
+	s.UnmarshalTOML([]byte(`'1'`))
+	assert.Equal(t, int64(1), s.Size)
+
+	s = Size{}
+	s.UnmarshalTOML([]byte(`"1GB"`))
+	assert.Equal(t, int64(1000*1000*1000), s.Size)
+
+	s = Size{}
+	s.UnmarshalTOML([]byte(`"12GiB"`))
+	assert.Equal(t, int64(12*1024*1024*1024), s.Size)
+}
+
 func TestCompressWithGzip(t *testing.T) {
 	testData := "the quick brown fox jumps over the lazy dog"
 	inputBuffer := bytes.NewBuffer([]byte(testData))