app/vmagent: add -remoteWrite.roundDigits command-line option for limiting the number of digits after the point for stored values

This commit also adds --vm-round-digits command-line option to vmctl tool.
This commit is contained in:
Aliaksandr Valialkin 2021-02-01 14:27:05 +02:00
parent 29a7067827
commit b2aa80e74b
13 changed files with 252 additions and 63 deletions

View file

@ -8,7 +8,6 @@ import (
"io/ioutil" "io/ioutil"
"net/http" "net/http"
"net/url" "net/url"
"strconv"
"strings" "strings"
"sync" "sync"
"time" "time"
@ -22,7 +21,7 @@ import (
) )
var ( var (
rateLimit = flagutil.NewArray("remoteWrite.rateLimit", "Optional rate limit in bytes per second for data sent to -remoteWrite.url. "+ rateLimit = flagutil.NewArrayInt("remoteWrite.rateLimit", "Optional rate limit in bytes per second for data sent to -remoteWrite.url. "+
"By default the rate limit is disabled. It can be useful for limiting load on remote storage when big amounts of buffered data "+ "By default the rate limit is disabled. It can be useful for limiting load on remote storage when big amounts of buffered data "+
"is sent after temporary unavailability of the remote storage") "is sent after temporary unavailability of the remote storage")
sendTimeout = flagutil.NewArrayDuration("remoteWrite.sendTimeout", "Timeout for sending a single block of data to -remoteWrite.url") sendTimeout = flagutil.NewArrayDuration("remoteWrite.sendTimeout", "Timeout for sending a single block of data to -remoteWrite.url")
@ -120,15 +119,9 @@ func newClient(argIdx int, remoteWriteURL, sanitizedURL string, fq *persistentqu
}, },
stopCh: make(chan struct{}), stopCh: make(chan struct{}),
} }
if bytesPerSec := rateLimit.GetOptionalArg(argIdx); bytesPerSec != "" { if bytesPerSec := rateLimit.GetOptionalArgOrDefault(argIdx, 0); bytesPerSec > 0 {
limit, err := strconv.ParseInt(bytesPerSec, 10, 64) logger.Infof("applying %d bytes per second rate limit for -remoteWrite.url=%q", bytesPerSec, sanitizedURL)
if err != nil { c.rl.perSecondLimit = int64(bytesPerSec)
logger.Fatalf("cannot parse -remoteWrite.rateLimit=%q for -remoteWrite.url=%q: %s", bytesPerSec, sanitizedURL, err)
}
if limit > 0 {
logger.Infof("applying %d bytes per second rate limit for -remoteWrite.url=%q", limit, sanitizedURL)
c.rl.perSecondLimit = limit
}
} }
c.rl.limitReached = metrics.GetOrCreateCounter(fmt.Sprintf(`vmagent_remote_write_rate_limit_reached_total{url=%q}`, c.sanitizedURL)) c.rl.limitReached = metrics.GetOrCreateCounter(fmt.Sprintf(`vmagent_remote_write_rate_limit_reached_total{url=%q}`, c.sanitizedURL))

View file

@ -7,6 +7,7 @@ import (
"time" "time"
"github.com/VictoriaMetrics/VictoriaMetrics/lib/bytesutil" "github.com/VictoriaMetrics/VictoriaMetrics/lib/bytesutil"
"github.com/VictoriaMetrics/VictoriaMetrics/lib/decimal"
"github.com/VictoriaMetrics/VictoriaMetrics/lib/fasttime" "github.com/VictoriaMetrics/VictoriaMetrics/lib/fasttime"
"github.com/VictoriaMetrics/VictoriaMetrics/lib/flagutil" "github.com/VictoriaMetrics/VictoriaMetrics/lib/flagutil"
"github.com/VictoriaMetrics/VictoriaMetrics/lib/persistentqueue" "github.com/VictoriaMetrics/VictoriaMetrics/lib/persistentqueue"
@ -35,9 +36,11 @@ type pendingSeries struct {
periodicFlusherWG sync.WaitGroup periodicFlusherWG sync.WaitGroup
} }
func newPendingSeries(pushBlock func(block []byte)) *pendingSeries { func newPendingSeries(pushBlock func(block []byte), significantFigures, roundDigits int) *pendingSeries {
var ps pendingSeries var ps pendingSeries
ps.wr.pushBlock = pushBlock ps.wr.pushBlock = pushBlock
ps.wr.significantFigures = significantFigures
ps.wr.roundDigits = roundDigits
ps.stopCh = make(chan struct{}) ps.stopCh = make(chan struct{})
ps.periodicFlusherWG.Add(1) ps.periodicFlusherWG.Add(1)
go func() { go func() {
@ -85,9 +88,17 @@ type writeRequest struct {
// Move lastFlushTime to the top of the struct in order to guarantee atomic access on 32-bit architectures. // Move lastFlushTime to the top of the struct in order to guarantee atomic access on 32-bit architectures.
lastFlushTime uint64 lastFlushTime uint64
wr prompbmarshal.WriteRequest // pushBlock is called when whe write request is ready to be sent.
pushBlock func(block []byte) pushBlock func(block []byte)
// How many significant figures must be left before sending the writeRequest to pushBlock.
significantFigures int
// How many decimal digits after point must be left before sending the writeRequest to pushBlock.
roundDigits int
wr prompbmarshal.WriteRequest
tss []prompbmarshal.TimeSeries tss []prompbmarshal.TimeSeries
labels []prompbmarshal.Label labels []prompbmarshal.Label
@ -96,6 +107,8 @@ type writeRequest struct {
} }
func (wr *writeRequest) reset() { func (wr *writeRequest) reset() {
// Do not reset pushBlock, significantFigures and roundDigits, since they are re-used.
wr.wr.Timeseries = nil wr.wr.Timeseries = nil
for i := range wr.tss { for i := range wr.tss {
@ -114,11 +127,28 @@ func (wr *writeRequest) reset() {
func (wr *writeRequest) flush() { func (wr *writeRequest) flush() {
wr.wr.Timeseries = wr.tss wr.wr.Timeseries = wr.tss
wr.adjustSampleValues()
atomic.StoreUint64(&wr.lastFlushTime, fasttime.UnixTimestamp()) atomic.StoreUint64(&wr.lastFlushTime, fasttime.UnixTimestamp())
pushWriteRequest(&wr.wr, wr.pushBlock) pushWriteRequest(&wr.wr, wr.pushBlock)
wr.reset() wr.reset()
} }
func (wr *writeRequest) adjustSampleValues() {
samples := wr.samples
if n := wr.significantFigures; n > 0 {
for i := range samples {
s := &samples[i]
s.Value = decimal.RoundToSignificantFigures(s.Value, n)
}
}
if n := wr.roundDigits; n < 100 {
for i := range samples {
s := &samples[i]
s.Value = decimal.RoundToDecimalDigits(s.Value, n)
}
}
}
func (wr *writeRequest) push(src []prompbmarshal.TimeSeries) { func (wr *writeRequest) push(src []prompbmarshal.TimeSeries) {
tssDst := wr.tss tssDst := wr.tss
for i := range src { for i := range src {

View file

@ -7,7 +7,6 @@ import (
"sync/atomic" "sync/atomic"
"github.com/VictoriaMetrics/VictoriaMetrics/lib/cgroup" "github.com/VictoriaMetrics/VictoriaMetrics/lib/cgroup"
"github.com/VictoriaMetrics/VictoriaMetrics/lib/decimal"
"github.com/VictoriaMetrics/VictoriaMetrics/lib/flagutil" "github.com/VictoriaMetrics/VictoriaMetrics/lib/flagutil"
"github.com/VictoriaMetrics/VictoriaMetrics/lib/logger" "github.com/VictoriaMetrics/VictoriaMetrics/lib/logger"
"github.com/VictoriaMetrics/VictoriaMetrics/lib/memory" "github.com/VictoriaMetrics/VictoriaMetrics/lib/memory"
@ -31,9 +30,13 @@ var (
"for each -remoteWrite.url. When buffer size reaches the configured maximum, then old data is dropped when adding new data to the buffer. "+ "for each -remoteWrite.url. When buffer size reaches the configured maximum, then old data is dropped when adding new data to the buffer. "+
"Buffered data is stored in ~500MB chunks, so the minimum practical value for this flag is 500000000. "+ "Buffered data is stored in ~500MB chunks, so the minimum practical value for this flag is 500000000. "+
"Disk usage is unlimited if the value is set to 0") "Disk usage is unlimited if the value is set to 0")
significantFigures = flag.Int("remoteWrite.significantFigures", 0, "The number of significant figures to leave in metric values before writing them to remote storage. "+ significantFigures = flagutil.NewArrayInt("remoteWrite.significantFigures", "The number of significant figures to leave in metric values before writing them "+
"See https://en.wikipedia.org/wiki/Significant_figures . Zero value saves all the significant figures. "+ "to remote storage. See https://en.wikipedia.org/wiki/Significant_figures . Zero value saves all the significant figures. "+
"This option may be used for increasing on-disk compression level for the stored metrics") "This option may be used for improving data compression for the stored metrics. See also -remoteWrite.roundDigits")
roundDigits = flagutil.NewArrayInt("remoteWrite.roundDigits", "Round metric values to this number of decimal digits after the point before writing them to remote storage. "+
"Examples: -remoteWrite.roundDigits=2 would round 1.236 to 1.24, while -remoteWrite.roundDigits=-1 would round 126.78 to 130. "+
"By default digits rounding is disabled. Set it to 100 for disabling it for a particular remote storage. "+
"This option may be used for improving data compression for the stored metrics")
) )
var rwctxs []*remoteWriteCtx var rwctxs []*remoteWriteCtx
@ -137,17 +140,6 @@ func Stop() {
// //
// Note that wr may be modified by Push due to relabeling and rounding. // Note that wr may be modified by Push due to relabeling and rounding.
func Push(wr *prompbmarshal.WriteRequest) { func Push(wr *prompbmarshal.WriteRequest) {
if *significantFigures > 0 {
// Round values according to significantFigures
for i := range wr.Timeseries {
samples := wr.Timeseries[i].Samples
for j := range samples {
s := &samples[j]
s.Value = decimal.Round(s.Value, *significantFigures)
}
}
}
var rctx *relabelCtx var rctx *relabelCtx
rcs := allRelabelConfigs.Load().(*relabelConfigs) rcs := allRelabelConfigs.Load().(*relabelConfigs)
prcsGlobal := rcs.global prcsGlobal := rcs.global
@ -213,9 +205,11 @@ func newRemoteWriteCtx(argIdx int, remoteWriteURL string, maxInmemoryBlocks int,
return float64(fq.GetInmemoryQueueLen()) return float64(fq.GetInmemoryQueueLen())
}) })
c := newClient(argIdx, remoteWriteURL, sanitizedURL, fq, *queues) c := newClient(argIdx, remoteWriteURL, sanitizedURL, fq, *queues)
sf := significantFigures.GetOptionalArgOrDefault(argIdx, 0)
rd := roundDigits.GetOptionalArgOrDefault(argIdx, 100)
pss := make([]*pendingSeries, *queues) pss := make([]*pendingSeries, *queues)
for i := range pss { for i := range pss {
pss[i] = newPendingSeries(fq.MustWriteBlock) pss[i] = newPendingSeries(fq.MustWriteBlock, sf, rd)
} }
return &remoteWriteCtx{ return &remoteWriteCtx{
idx: argIdx, idx: argIdx,

View file

@ -413,15 +413,20 @@ Moreover, such values may be just a result of [floating point arithmetic](https:
create a [false precision](https://en.wikipedia.org/wiki/False_precision) and result into bad compression ratio create a [false precision](https://en.wikipedia.org/wiki/False_precision) and result into bad compression ratio
according to [information theory](https://en.wikipedia.org/wiki/Information_theory). according to [information theory](https://en.wikipedia.org/wiki/Information_theory).
The `--vm-significant-figures` flag allows to limit the number of significant figures. It takes no effect if set `vmctl` provides the following flags for improving data compression:
to 0 (by default), but set `--vm-significant-figures=5` and `102.342305` will be rounded to `102.34`. Such value will
have much higher compression ratio comparing to previous one and will save some extra disk space after the migration. * `--vm-round-digits` flag for rounding processed values to the given number of decimal digits after the point.
The most common case for using this flag is to reduce number of significant figures for time series storing aggregation For example, `--vm-round-digits=2` would round `1.2345` to `1.23`. By default the rounding is disabled.
results such as `average`, `rate`, etc.
* `--vm-significant-figures` flag for limiting the number of significant figures in processed values. It takes no effect if set
to 0 (by default), but set `--vm-significant-figures=5` and `102.342305` will be rounded to `102.34`.
The most common case for using these flags is to improve data compression for time series storing aggregation
results such as `average`, `rate`, etc.
### Adding extra labels ### Adding extra labels
`vmctl` allows to add extra labels to all imported series. It can be achived with flag `--vm-extra-label label=value`. `vmctl` allows to add extra labels to all imported series. It can be achived with flag `--vm-extra-label label=value`.
If multiple labels needs to be added, set flag for each label, for example, `--vm-extra-label label1=value1 --vm-extra-label label2=value2`. If multiple labels needs to be added, set flag for each label, for example, `--vm-extra-label label1=value1 --vm-extra-label label2=value2`.
If timeseries already have label, that must be added with `--vm-extra-label` flag, flag has priority and will override label value from timeseries. If timeseries already have label, that must be added with `--vm-extra-label` flag, flag has priority and will override label value from timeseries.

View file

@ -29,6 +29,7 @@ const (
vmCompress = "vm-compress" vmCompress = "vm-compress"
vmBatchSize = "vm-batch-size" vmBatchSize = "vm-batch-size"
vmSignificantFigures = "vm-significant-figures" vmSignificantFigures = "vm-significant-figures"
vmRoundDigits = "vm-round-digits"
vmExtraLabel = "vm-extra-label" vmExtraLabel = "vm-extra-label"
) )
@ -77,6 +78,13 @@ var (
Value: 0, Value: 0,
Usage: "The number of significant figures to leave in metric values before importing. " + Usage: "The number of significant figures to leave in metric values before importing. " +
"See https://en.wikipedia.org/wiki/Significant_figures. Zero value saves all the significant figures. " + "See https://en.wikipedia.org/wiki/Significant_figures. Zero value saves all the significant figures. " +
"This option may be used for increasing on-disk compression level for the stored metrics. " +
"See also --vm-round-digits option",
},
&cli.IntFlag{
Name: vmRoundDigits,
Value: 100,
Usage: "Round metric values to the given number of decimal digits after the point. " +
"This option may be used for increasing on-disk compression level for the stored metrics", "This option may be used for increasing on-disk compression level for the stored metrics",
}, },
&cli.StringSliceFlag{ &cli.StringSliceFlag{

View file

@ -153,6 +153,7 @@ func initConfigVM(c *cli.Context) vm.Config {
AccountID: c.String(vmAccountID), AccountID: c.String(vmAccountID),
BatchSize: c.Int(vmBatchSize), BatchSize: c.Int(vmBatchSize),
SignificantFigures: c.Int(vmSignificantFigures), SignificantFigures: c.Int(vmSignificantFigures),
RoundDigits: c.Int(vmRoundDigits),
ExtraLabels: c.StringSlice(vmExtraLabel), ExtraLabels: c.StringSlice(vmExtraLabel),
} }
} }

View file

@ -42,6 +42,9 @@ type Config struct {
// in metric values before importing. // in metric values before importing.
// Zero value saves all the significant decimal places // Zero value saves all the significant decimal places
SignificantFigures int SignificantFigures int
// RoundDigits defines the number of decimal digits after the point that must be left
// in metric values before importing.
RoundDigits int
// ExtraLabels that will be added to all imported series. Must be in label=value format. // ExtraLabels that will be added to all imported series. Must be in label=value format.
ExtraLabels []string ExtraLabels []string
} }
@ -136,7 +139,7 @@ func NewImporter(cfg Config) (*Importer, error) {
for i := 0; i < int(cfg.Concurrency); i++ { for i := 0; i < int(cfg.Concurrency); i++ {
go func() { go func() {
defer im.wg.Done() defer im.wg.Done()
im.startWorker(cfg.BatchSize, cfg.SignificantFigures) im.startWorker(cfg.BatchSize, cfg.SignificantFigures, cfg.RoundDigits)
}() }()
} }
im.ResetStats() im.ResetStats()
@ -170,7 +173,7 @@ func (im *Importer) Close() {
}) })
} }
func (im *Importer) startWorker(batchSize, significantFigures int) { func (im *Importer) startWorker(batchSize, significantFigures, roundDigits int) {
var batch []*TimeSeries var batch []*TimeSeries
var dataPoints int var dataPoints int
var waitForBatch time.Time var waitForBatch time.Time
@ -192,9 +195,13 @@ func (im *Importer) startWorker(batchSize, significantFigures int) {
} }
if significantFigures > 0 { if significantFigures > 0 {
// Round values according to significantFigures
for i, v := range ts.Values { for i, v := range ts.Values {
ts.Values[i] = decimal.Round(v, significantFigures) ts.Values[i] = decimal.RoundToSignificantFigures(v, significantFigures)
}
}
if roundDigits < 100 {
for i, v := range ts.Values {
ts.Values[i] = decimal.RoundToDecimalDigits(v, roundDigits)
} }
} }

View file

@ -6,6 +6,7 @@
* FEATURE: added `-loggerTimezone` command-line flag for adjusting time zone for timestamps in log messages. By default UTC is used. * FEATURE: added `-loggerTimezone` command-line flag for adjusting time zone for timestamps in log messages. By default UTC is used.
* FEATURE: added `-search.maxStepForPointsAdjustment` command-line flag, which can be used for disabling adjustment for points returned by `/api/v1/query_range` handler if such points have timestamps closer than `-search.latencyOffset` to the current time. Such points may contain incomplete data, so they are substituted by the previous values for `step` query args smaller than one minute by default. * FEATURE: added `-search.maxStepForPointsAdjustment` command-line flag, which can be used for disabling adjustment for points returned by `/api/v1/query_range` handler if such points have timestamps closer than `-search.latencyOffset` to the current time. Such points may contain incomplete data, so they are substituted by the previous values for `step` query args smaller than one minute by default.
* FEATURE: vmalert: added `-datasource.queryStep` command-line flag for passing optional `step` query arg to `/api/v1/query` endpoint. See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/1025 * FEATURE: vmalert: added `-datasource.queryStep` command-line flag for passing optional `step` query arg to `/api/v1/query` endpoint. See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/1025
* FEATURE: vmagent: added `-remoteWrite.roundDigits` command-line option for rounding metric values to the given number of decimal digits after the point before sending the metric to the corresponding `-remoteWrite.url`. This option can be used for improving data compression on the remote storage, because values with lower number of decimal digits can be compressed better than values with bigger number of decimal digits.
* FEATURE: vmagent: added `-remoteWrite.rateLimit` command-line flag for limiting data transfer rate to `-remoteWrite.url`. This may be useful when big amounts of buffered data is sent after temporarily unavailability of the remote storage. See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/1035 * FEATURE: vmagent: added `-remoteWrite.rateLimit` command-line flag for limiting data transfer rate to `-remoteWrite.url`. This may be useful when big amounts of buffered data is sent after temporarily unavailability of the remote storage. See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/1035
* FEATURE: vmagent: export `vm_promscrape_scrapes_failed_per_url_total` and `vm_promscrape_scrapes_skipped_by_sample_limit_per_url_total` counters, which may help identifying improperly working scrape targets. * FEATURE: vmagent: export `vm_promscrape_scrapes_failed_per_url_total` and `vm_promscrape_scrapes_skipped_by_sample_limit_per_url_total` counters, which may help identifying improperly working scrape targets.

View file

@ -413,15 +413,20 @@ Moreover, such values may be just a result of [floating point arithmetic](https:
create a [false precision](https://en.wikipedia.org/wiki/False_precision) and result into bad compression ratio create a [false precision](https://en.wikipedia.org/wiki/False_precision) and result into bad compression ratio
according to [information theory](https://en.wikipedia.org/wiki/Information_theory). according to [information theory](https://en.wikipedia.org/wiki/Information_theory).
The `--vm-significant-figures` flag allows to limit the number of significant figures. It takes no effect if set `vmctl` provides the following flags for improving data compression:
to 0 (by default), but set `--vm-significant-figures=5` and `102.342305` will be rounded to `102.34`. Such value will
have much higher compression ratio comparing to previous one and will save some extra disk space after the migration. * `--vm-round-digits` flag for rounding processed values to the given number of decimal digits after the point.
The most common case for using this flag is to reduce number of significant figures for time series storing aggregation For example, `--vm-round-digits=2` would round `1.2345` to `1.23`. By default the rounding is disabled.
results such as `average`, `rate`, etc.
* `--vm-significant-figures` flag for limiting the number of significant figures in processed values. It takes no effect if set
to 0 (by default), but set `--vm-significant-figures=5` and `102.342305` will be rounded to `102.34`.
The most common case for using these flags is to improve data compression for time series storing aggregation
results such as `average`, `rate`, etc.
### Adding extra labels ### Adding extra labels
`vmctl` allows to add extra labels to all imported series. It can be achived with flag `--vm-extra-label label=value`. `vmctl` allows to add extra labels to all imported series. It can be achived with flag `--vm-extra-label label=value`.
If multiple labels needs to be added, set flag for each label, for example, `--vm-extra-label label1=value1 --vm-extra-label label2=value2`. If multiple labels needs to be added, set flag for each label, for example, `--vm-extra-label label1=value1 --vm-extra-label label2=value2`.
If timeseries already have label, that must be added with `--vm-extra-label` flag, flag has priority and will override label value from timeseries. If timeseries already have label, that must be added with `--vm-extra-label` flag, flag has priority and will override label value from timeseries.

View file

@ -298,8 +298,21 @@ func maxUpExponent(v int64) int16 {
} }
} }
// Round f to value with the given number of significant figures. // RoundToDecimalDigits rounds f to the given number of decimal digits after the point.
func Round(f float64, digits int) float64 { //
// See also RoundToSignificantFigures.
func RoundToDecimalDigits(f float64, digits int) float64 {
if digits <= -100 || digits >= 100 {
return f
}
m := math.Pow10(digits)
return math.Round(f*m) / m
}
// RoundToSignificantFigures rounds f to value with the given number of significant figures.
//
// See also RoundToDecimalDigits.
func RoundToSignificantFigures(f float64, digits int) float64 {
if digits <= 0 || digits >= 18 { if digits <= 0 || digits >= 18 {
return f return f
} }

View file

@ -7,10 +7,34 @@ import (
"testing" "testing"
) )
func TestRoundInplace(t *testing.T) { func TestRoundToDecimalDigits(t *testing.T) {
f := func(f float64, digits int, resultExpected float64) { f := func(f float64, digits int, resultExpected float64) {
t.Helper() t.Helper()
result := Round(f, digits) result := RoundToDecimalDigits(f, digits)
if math.IsNaN(result) {
if !math.IsNaN(resultExpected) {
t.Fatalf("unexpected result; got %v; want %v", result, resultExpected)
}
}
if result != resultExpected {
t.Fatalf("unexpected result; got %v; want %v", result, resultExpected)
}
}
f(12.34, 0, 12)
f(12.57, 0, 13)
f(-1.578, 2, -1.58)
f(-1.578, 3, -1.578)
f(1234, -2, 1200)
f(1235, -1, 1240)
f(1234, 0, 1234)
f(1234.6, 0, 1235)
f(123.4e-99, 99, 123e-99)
}
func TestRoundToSignificantFigures(t *testing.T) {
f := func(f float64, digits int, resultExpected float64) {
t.Helper()
result := RoundToSignificantFigures(f, digits)
if math.IsNaN(result) { if math.IsNaN(result) {
if !math.IsNaN(resultExpected) { if !math.IsNaN(resultExpected) {
t.Fatalf("unexpected result; got %v; want %v", result, resultExpected) t.Fatalf("unexpected result; got %v; want %v", result, resultExpected)

View file

@ -35,6 +35,15 @@ func NewArrayBool(name, description string) *ArrayBool {
return &a return &a
} }
// NewArrayInt returns new ArrayInt with the given name and description.
func NewArrayInt(name string, description string) *ArrayInt {
description += "\nSupports `array` of values separated by comma" +
" or specified via multiple flags."
var a ArrayInt
flag.Var(&a, name, description)
return &a
}
// Array is a flag that holds an array of values. // Array is a flag that holds an array of values.
// //
// It may be set either by specifying multiple flags with the given name // It may be set either by specifying multiple flags with the given name
@ -223,3 +232,41 @@ func (a *ArrayDuration) GetOptionalArgOrDefault(argIdx int, defaultValue time.Du
} }
return x[argIdx] return x[argIdx]
} }
// ArrayInt is flag that holds an array of ints.
type ArrayInt []int
// String implements flag.Value interface
func (a *ArrayInt) String() string {
x := *a
formattedInts := make([]string, len(x))
for i, v := range x {
formattedInts[i] = strconv.Itoa(v)
}
return strings.Join(formattedInts, ",")
}
// Set implements flag.Value interface
func (a *ArrayInt) Set(value string) error {
values := parseArrayValues(value)
for _, v := range values {
n, err := strconv.Atoi(v)
if err != nil {
return err
}
*a = append(*a, n)
}
return nil
}
// GetOptionalArg returns optional arg under the given argIdx.
func (a *ArrayInt) GetOptionalArgOrDefault(argIdx int, defaultValue int) int {
x := *a
if argIdx < len(x) {
return x[argIdx]
}
if len(x) == 1 {
return x[0]
}
return defaultValue
}

View file

@ -12,14 +12,18 @@ var (
fooFlag Array fooFlag Array
fooFlagDuration ArrayDuration fooFlagDuration ArrayDuration
fooFlagBool ArrayBool fooFlagBool ArrayBool
fooFlagInt ArrayInt
) )
func init() { func init() {
os.Args = append(os.Args, "--fooFlag=foo", "--fooFlag=bar", "--fooFlagDuration=10s", "--fooFlagDuration=5m") os.Args = append(os.Args, "--fooFlag=foo", "--fooFlag=bar")
os.Args = append(os.Args, "--fooFlagDuration=10s", "--fooFlagDuration=5m")
os.Args = append(os.Args, "--fooFlagBool=true", "--fooFlagBool=false,true", "--fooFlagBool") os.Args = append(os.Args, "--fooFlagBool=true", "--fooFlagBool=false,true", "--fooFlagBool")
os.Args = append(os.Args, "--fooFlagInt=1", "--fooFlagInt=2,3")
flag.Var(&fooFlag, "fooFlag", "test") flag.Var(&fooFlag, "fooFlag", "test")
flag.Var(&fooFlagDuration, "fooFlagDuration", "test") flag.Var(&fooFlagDuration, "fooFlagDuration", "test")
flag.Var(&fooFlagBool, "fooFlagBool", "test") flag.Var(&fooFlagBool, "fooFlagBool", "test")
flag.Var(&fooFlagInt, "fooFlagInt", "test")
} }
func TestMain(m *testing.M) { func TestMain(m *testing.M) {
@ -28,16 +32,16 @@ func TestMain(m *testing.M) {
} }
func TestArray(t *testing.T) { func TestArray(t *testing.T) {
expected := map[string]struct{}{ expected := []string{
"foo": {}, "foo",
"bar": {}, "bar",
} }
if len(expected) != len(fooFlag) { if len(expected) != len(fooFlag) {
t.Errorf("len array flag (%d) is not equal to %d", len(fooFlag), len(expected)) t.Errorf("len array flag (%d) is not equal to %d", len(fooFlag), len(expected))
} }
for _, i := range fooFlag { for i, v := range fooFlag {
if _, ok := expected[i]; !ok { if v != expected[i] {
t.Errorf("unexpected item in array %v", i) t.Errorf("unexpected item in array %q", v)
} }
} }
} }
@ -101,16 +105,16 @@ func TestArrayString(t *testing.T) {
} }
func TestArrayDuration(t *testing.T) { func TestArrayDuration(t *testing.T) {
expected := map[time.Duration]struct{}{ expected := []time.Duration{
time.Second * 10: {}, time.Second * 10,
time.Minute * 5: {}, time.Minute * 5,
} }
if len(expected) != len(fooFlagDuration) { if len(expected) != len(fooFlagDuration) {
t.Errorf("len array flag (%d) is not equal to %d", len(fooFlag), len(expected)) t.Errorf("len array flag (%d) is not equal to %d", len(fooFlag), len(expected))
} }
for _, i := range fooFlagDuration { for i, v := range fooFlagDuration {
if _, ok := expected[i]; !ok { if v != expected[i] {
t.Errorf("unexpected item in array %v", i) t.Errorf("unexpected item in array %s", v)
} }
} }
} }
@ -130,7 +134,7 @@ func TestArrayDurationSet(t *testing.T) {
} }
func TestArrayDurationGetOptionalArg(t *testing.T) { func TestArrayDurationGetOptionalArg(t *testing.T) {
f := func(s string, argIdx int, expectedValue time.Duration, defaultValue time.Duration) { f := func(s string, argIdx int, expectedValue, defaultValue time.Duration) {
t.Helper() t.Helper()
var a ArrayDuration var a ArrayDuration
_ = a.Set(s) _ = a.Set(s)
@ -219,3 +223,60 @@ func TestArrayBoolString(t *testing.T) {
f("true,false") f("true,false")
f("false,true") f("false,true")
} }
func TestArrayInt(t *testing.T) {
expected := []int{1, 2, 3}
if len(expected) != len(fooFlagInt) {
t.Errorf("len array flag (%d) is not equal to %d", len(fooFlag), len(expected))
}
for i, n := range fooFlagInt {
if n != expected[i] {
t.Errorf("unexpected item in array %d", n)
}
}
}
func TestArrayIntSet(t *testing.T) {
f := func(s string, expectedValues []int) {
t.Helper()
var a ArrayInt
_ = a.Set(s)
if !reflect.DeepEqual([]int(a), expectedValues) {
t.Fatalf("unexpected values parsed;\ngot\n%q\nwant\n%q", a, expectedValues)
}
}
f("", nil)
f(`1`, []int{1})
f(`-2,3,-64`, []int{-2, 3, -64})
}
func TestArrayIntGetOptionalArg(t *testing.T) {
f := func(s string, argIdx int, expectedValue, defaultValue int) {
t.Helper()
var a ArrayInt
_ = a.Set(s)
v := a.GetOptionalArgOrDefault(argIdx, defaultValue)
if v != expectedValue {
t.Fatalf("unexpected value; got %d; want %d", v, expectedValue)
}
}
f("", 0, 123, 123)
f("", 1, -34, -34)
f("10,1", 1, 1, 234)
f("10", 3, 10, -34)
}
func TestArrayIntString(t *testing.T) {
f := func(s string) {
t.Helper()
var a ArrayInt
_ = a.Set(s)
result := a.String()
if result != s {
t.Fatalf("unexpected string;\ngot\n%s\nwant\n%s", result, s)
}
}
f("")
f("10,1")
f("-5,1,123")
}