Add a format for object-size providing a wieighted random based on configured buckets

cli/generator.go

@@ -27,6 +27,8 @@ import (
 
 	"github.com/minio/cli"
 	"github.com/minio/warp/pkg/generator"
+
+	hist "github.com/jfsmig/prng/histogram"
 )
 
 var genFlags = []cli.Flag{
@@ -83,28 +85,38 @@ func newGenSource(ctx *cli.Context, sizeField string) func() generator.Source {
 		generator.WithCustomPrefix(ctx.String("prefix")),
 		generator.WithPrefixSize(prefixSize),
 	}
-	tokens := strings.Split(ctx.String(sizeField), ",")
-	switch len(tokens) {
-	case 1:
-		size, err := toSize(tokens[0])
-		if err != nil {
-			fatalIf(probe.NewError(err), "Invalid obj.size specified")
-		}
-		opts = append(opts, generator.WithSize(int64(size)))
-	case 2:
-		minSize, err := toSize(tokens[0])
-		if err != nil {
-			fatalIf(probe.NewError(err), "Invalid min obj.size specified")
+	if strings.IndexRune(ctx.String(sizeField), ':') > 0 {
+		if _, err := hist.ParseCSV(ctx.String(sizeField)); err != nil {
+			fatalIf(probe.NewError(err), "Invalid histogram format for the size parameter")
+		} else {
+			opts = append(opts, generator.WithSizeHistograms(ctx.String(sizeField)))
 		}
-		maxSize, err := toSize(tokens[1])
-		if err != nil {
-			fatalIf(probe.NewError(err), "Invalid max obj.size specified")
+	} else {
+		tokens := strings.Split(ctx.String(sizeField), ",")
+		switch len(tokens) {
+		case 1:
+			size, err := toSize(tokens[0])
+			if err != nil {
+				fatalIf(probe.NewError(err), "Invalid obj.size specified")
+			}
+			opts = append(opts, generator.WithSize(int64(size)))
+		case 2:
+			minSize, err := toSize(tokens[0])
+			if err != nil {
+				fatalIf(probe.NewError(err), "Invalid min obj.size specified")
+			}
+			maxSize, err := toSize(tokens[1])
+			if err != nil {
+				fatalIf(probe.NewError(err), "Invalid max obj.size specified")
+			}
+			opts = append(opts, generator.WithMinMaxSize(int64(minSize), int64(maxSize)))
+		default:
+			fatalIf(probe.NewError(fmt.Errorf("unexpected obj.size specified: %s", ctx.String(sizeField))), "Invalid obj.size parameter")
 		}
-		opts = append(opts, generator.WithMinMaxSize(int64(minSize), int64(maxSize)))
-	default:
-		fatalIf(probe.NewError(fmt.Errorf("unexpected obj.size specified: %s", ctx.String(sizeField))), "Invalid obj.size parameter")
+
+		opts = append([]generator.Option{g.Apply()}, append(opts, generator.WithRandomSize(ctx.Bool("obj.randsize")))...)
 	}
-	opts = append([]generator.Option{g.Apply()}, append(opts, generator.WithRandomSize(ctx.Bool("obj.randsize")))...)
+
 	src, err := generator.NewFn(opts...)
 	fatalIf(probe.NewError(err), "Unable to create data generator")
 	return src

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/dustin/go-humanize v1.0.1
 	github.com/fatih/color v1.17.0
 	github.com/influxdata/influxdb-client-go/v2 v2.13.0
+	github.com/jfsmig/prng v0.0.2
 	github.com/klauspost/compress v1.17.9
 	github.com/minio/cli v1.24.2
 	github.com/minio/madmin-go/v3 v3.0.51

go.sum

@@ -48,6 +48,8 @@ github.com/influxdata/influxdb-client-go/v2 v2.13.0 h1:ioBbLmR5NMbAjP4UVA5r9b5xG
 github.com/influxdata/influxdb-client-go/v2 v2.13.0/go.mod h1:k+spCbt9hcvqvUiz0sr5D8LolXHqAAOfPw9v/RIRHl4=
 github.com/influxdata/line-protocol v0.0.0-20210922203350-b1ad95c89adf h1:7JTmneyiNEwVBOHSjoMxiWAqB992atOeepeFYegn5RU=
 github.com/influxdata/line-protocol v0.0.0-20210922203350-b1ad95c89adf/go.mod h1:xaLFMmpvUxqXtVkUJfg9QmT88cDaCJ3ZKgdZ78oO8Qo=
+github.com/jfsmig/prng v0.0.2 h1:aZun+YgmBnUyhqvI+EDjwmOYc1kCPsihdEr9V/1YlGA=
+github.com/jfsmig/prng v0.0.2/go.mod h1:bz1fX1aizp8/Lu1thLzfirh5uExjC1lVwB8SSt6ExpE=
 github.com/juju/gnuflag v0.0.0-20171113085948-2ce1bb71843d/go.mod h1:2PavIy+JPciBPrBUjwbNvtwB6RQlve+hkpll6QSNmOE=
 github.com/klauspost/compress v1.17.9 h1:6KIumPrER1LHsvBVuDa0r5xaG0Es51mhhB9BQB2qeMA=
 github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=

pkg/generator/options.go

@@ -20,6 +20,8 @@ package generator
 import (
 	"errors"
 	"math/rand"
+
+	hist "github.com/jfsmig/prng/histogram"
 )
 
 // Options provides options.
@@ -33,6 +35,10 @@ type Options struct {
 	totalSize    int64
 	randomPrefix int
 	randSize     bool
+
+	// Activates the use of a distribution of sizes
+	flagSizesDistribution bool
+	sizesDistribution     hist.Int64Distribution
 }
 
 // OptionApplier allows to abstract generator options.
@@ -42,6 +48,9 @@ type OptionApplier interface {
 
 // getSize will return a size for an object.
 func (o Options) getSize(rng *rand.Rand) int64 {
+	if o.flagSizesDistribution {
+		return o.sizesDistribution.Poll(rng)
+	}
 	if !o.randSize {
 		return o.totalSize
 	}
@@ -59,20 +68,32 @@ func defaultOptions() Options {
 	return o
 }
 
+func WithSizeHistograms(encoded string) Option {
+	return func(o *Options) error {
+		var err error
+		o.sizesDistribution, err = hist.ParseCSV(encoded)
+		if err != nil {
+			return err
+		}
+		o.flagSizesDistribution = true
+		return nil
+	}
+}
+
 // WithMinMaxSize sets the min and max size of the generated data.
 func WithMinMaxSize(min, max int64) Option {
 	return func(o *Options) error {
 		if min <= 0 {
-			return errors.New("WithSize: minSize must be >= 0")
+			return errors.New("WithMinMaxSize: minSize must be >= 0")
 		}
 		if max < 0 {
-			return errors.New("WithSize: maxSize must be > 0")
+			return errors.New("WithMinMaxSize: maxSize must be > 0")
 		}
 		if min > max {
-			return errors.New("WithSize: minSize must be < maxSize")
+			return errors.New("WithMinMaxSize: minSize must be < maxSize")
 		}
 		if o.randSize && max < 256 {
-			return errors.New("WithSize: random sized objects should be at least 256 bytes")
+			return errors.New("WithMinMaxSize: random sized objects should be at least 256 bytes")
 		}
 
 		o.totalSize = max