VictoriaMetrics/lib/flagutil/dict.go
Zakhar Bessarab efb7213f6b
lib/flagutil/dict: properly update default value in case there is no key value set ()
### Describe Your Changes

If a dict flag has only one value without a prefix it is supposed to
replace default value.

Previously, when flag was set to `-flag=2` and the default value in
`NewDictInt` was set to 1 the resulting value for any `flag.Get()` call
would be 1 which is not expected.

This commit updates default value for the flag in case there is only one
entry for flag and the entry is a number without a key.

This affects cluster version and specifically `replicationFactor` flag
usage with vmstorage [node
groups](https://docs.victoriametrics.com/cluster-victoriametrics/#vmstorage-groups-at-vmselect).
Previously, the following configuration would effectively be ignored:
```
/path/to/vmselect \
 -replicationFactor=2 \
 -storageNode=g1/host1,g1/host2,g1/host3 \
 -storageNode=g2/host4,g2/host5,g2/host6 \
 -storageNode=g3/host7,g3/host8,g3/host9
```

Changes from this PR will force default value for `replicationFactor`
flag to be set to `2` which is expected as the result of this
configuration.


---------

Signed-off-by: Zakhar Bessarab <z.bessarab@victoriametrics.com>
2024-10-17 11:11:45 -03:00

115 lines
2.5 KiB
Go

package flagutil
import (
"encoding/json"
"flag"
"fmt"
"strconv"
"strings"
)
// DictInt allows specifying a dictionary of named ints in the form `name1:value1,...,nameN:valueN`.
type DictInt struct {
defaultValue int
kvs []kIntValue
}
type kIntValue struct {
k string
v int
}
// NewDictInt creates DictInt with the given name, defaultValue and description.
func NewDictInt(name string, defaultValue int, description string) *DictInt {
description += fmt.Sprintf(" (default %d)", defaultValue)
description += "\nSupports an `array` of `key:value` entries separated by comma or specified via multiple flags."
di := &DictInt{
defaultValue: defaultValue,
}
flag.Var(di, name, description)
return di
}
// String implements flag.Value interface
func (di *DictInt) String() string {
kvs := di.kvs
if len(kvs) == 1 && kvs[0].k == "" {
// Short form - a single int value
return strconv.Itoa(kvs[0].v)
}
formattedResults := make([]string, len(kvs))
for i, kv := range kvs {
formattedResults[i] = fmt.Sprintf("%s:%d", kv.k, kv.v)
}
return strings.Join(formattedResults, ",")
}
// Set implements flag.Value interface
func (di *DictInt) Set(value string) error {
values := parseArrayValues(value)
if len(di.kvs) == 0 && len(values) == 1 && strings.IndexByte(values[0], ':') < 0 {
v, err := strconv.Atoi(values[0])
if err != nil {
return err
}
di.kvs = append(di.kvs, kIntValue{
v: v,
})
di.defaultValue = v
return nil
}
for _, x := range values {
n := strings.IndexByte(x, ':')
if n < 0 {
return fmt.Errorf("missing ':' in %q", x)
}
k := x[:n]
v, err := strconv.Atoi(x[n+1:])
if err != nil {
return fmt.Errorf("cannot parse value for key=%q: %w", k, err)
}
if di.contains(k) {
return fmt.Errorf("duplicate value for key=%q: %d", k, v)
}
di.kvs = append(di.kvs, kIntValue{
k: k,
v: v,
})
}
return nil
}
func (di *DictInt) contains(key string) bool {
for _, kv := range di.kvs {
if kv.k == key {
return true
}
}
return false
}
// Get returns value for the given key.
//
// Default value is returned if key isn't found in di.
func (di *DictInt) Get(key string) int {
for _, kv := range di.kvs {
if kv.k == key {
return kv.v
}
}
return di.defaultValue
}
// ParseJSONMap parses s, which must contain JSON map of {"k1":"v1",...,"kN":"vN"}
func ParseJSONMap(s string) (map[string]string, error) {
if s == "" {
// Special case
return nil, nil
}
var m map[string]string
if err := json.Unmarshal([]byte(s), &m); err != nil {
return nil, err
}
return m, nil
}