Skip to content

Avoid calling flag.Parse() twice. #1997

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Jan 21, 2020
61 changes: 50 additions & 11 deletions cmd/cortex/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,36 +25,54 @@ func init() {
prometheus.MustRegister(version.NewCollector("cortex"))
}

const configFileOption = "config.file"

var testMode = false

func main() {
var (
cfg cortex.Config
configFile = ""
eventSampleRate int
ballastBytes int
mutexProfileFraction int
)
flag.StringVar(&configFile, "config.file", "", "Configuration file to load.")
flag.IntVar(&eventSampleRate, "event.sample-rate", 0, "How often to sample observability events (0 = never).")
flag.IntVar(&ballastBytes, "mem-ballast-size-bytes", 0, "Size of memory ballast to allocate.")
flag.IntVar(&mutexProfileFraction, "debug.mutex-profile-fraction", 0, "Fraction at which mutex profile vents will be reported, 0 to disable")

if mutexProfileFraction > 0 {
runtime.SetMutexProfileFraction(mutexProfileFraction)
}
configFile := parseConfigFileParameter()

// This sets default values from flags to the config.
// It needs to be called before parsing the config file!
flagext.RegisterFlags(&cfg)
flag.Parse()

if configFile != "" {
if err := LoadConfig(configFile, &cfg); err != nil {
fmt.Printf("error loading config from %s: %v\n", configFile, err)
fmt.Fprintf(os.Stderr, "error loading config from %s: %v\n", configFile, err)
if testMode {
return
}
os.Exit(1)
}
}

// Parse a second time, as command line flags should take precedent over the config file.
// Ignore -config.file here, since it was already parsed, but it's still present on command line.
flagext.IgnoredFlag(flag.CommandLine, configFileOption, "Configuration file to load.")
flag.IntVar(&eventSampleRate, "event.sample-rate", 0, "How often to sample observability events (0 = never).")
flag.IntVar(&ballastBytes, "mem-ballast-size-bytes", 0, "Size of memory ballast to allocate.")
flag.IntVar(&mutexProfileFraction, "debug.mutex-profile-fraction", 0, "Fraction at which mutex profile vents will be reported, 0 to disable")

if testMode {
// Don't exit on error in test mode. Just parse parameters, dump config and stop.
flag.CommandLine.Init(flag.CommandLine.Name(), flag.ContinueOnError)
flag.Parse()
DumpYaml(&cfg)
return
}

flag.Parse()

if mutexProfileFraction > 0 {
runtime.SetMutexProfileFraction(mutexProfileFraction)
}

// Validate the config once both the config file has been loaded
// and CLI flags parsed.
err := cfg.Validate()
Expand Down Expand Up @@ -90,6 +108,18 @@ func main() {
util.CheckFatal("initializing cortex", err)
}

// Parse -config.file option via separate flag set, to avoid polluting default one and calling flag.Parse on it twice.
func parseConfigFileParameter() string {
var configFile = ""
// ignore errors and any output here. Any flag errors will be reported by main flag.Parse() call.
fs := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
fs.SetOutput(ioutil.Discard)
fs.StringVar(&configFile, configFileOption, "", "") // usage not used in this function.
_ = fs.Parse(os.Args[1:])

return configFile
}

// LoadConfig read YAML-formatted config from filename into cfg.
func LoadConfig(filename string, cfg *cortex.Config) error {
buf, err := ioutil.ReadFile(filename)
Expand All @@ -104,3 +134,12 @@ func LoadConfig(filename string, cfg *cortex.Config) error {

return nil
}

func DumpYaml(cfg *cortex.Config) {
out, err := yaml.Marshal(cfg)
if err != nil {
fmt.Fprintln(os.Stderr, err)
} else {
fmt.Printf("%s\n", out)
}
}
163 changes: 163 additions & 0 deletions cmd/cortex/main_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,163 @@
package main

import (
"bytes"
"flag"
"io"
"io/ioutil"
"os"
"sync"
"testing"

"github.com/stretchr/testify/require"
)

func TestFlagParsing(t *testing.T) {
for name, tc := range map[string]struct {
arguments []string
yaml string
stdoutMessage string // string that must be included in stdout
stderrMessage string // string that must be included in stderr
}{
"help": {
arguments: []string{"-h"},
stderrMessage: configFileOption,
},

// check that config file is used
"config with unknown target": {
yaml: "target: unknown",
stderrMessage: "unrecognised module name: unknown",
},

"argument with unknown target": {
arguments: []string{"-target=unknown"},
stderrMessage: "unrecognised module name: unknown",
},

"unknown flag": {
arguments: []string{"-unknown.flag"},
stderrMessage: "-unknown.flag",
},

"config with wrong argument override": {
yaml: "target: ingester",
arguments: []string{"-target=unknown"},
stderrMessage: "unrecognised module name: unknown",
},

"default values": {
stdoutMessage: "target: all\n",
},

"config": {
yaml: "target: ingester",
stdoutMessage: "target: ingester\n",
},

"config with arguments override": {
yaml: "target: ingester",
arguments: []string{"-target=distributor"},
stdoutMessage: "target: distributor\n",
},

// we cannot test the happy path, as cortex would then fully start
} {
t.Run(name, func(t *testing.T) {
testSingle(t, tc.arguments, tc.yaml, []byte(tc.stdoutMessage), []byte(tc.stderrMessage))
})
}
}

func testSingle(t *testing.T, arguments []string, yaml string, stdoutMessage, stderrMessage []byte) {
oldArgs, oldStdout, oldStderr, oldTestMode := os.Args, os.Stdout, os.Stderr, testMode
defer func() {
os.Stdout = oldStdout
os.Stderr = oldStderr
os.Args = oldArgs
testMode = oldTestMode
}()

if yaml != "" {
tempFile, err := ioutil.TempFile("", "test")
require.NoError(t, err)

defer func() {
require.NoError(t, tempFile.Close())
require.NoError(t, os.Remove(tempFile.Name()))
}()

_, err = tempFile.WriteString(yaml)
require.NoError(t, err)

arguments = append([]string{"-" + configFileOption, tempFile.Name()}, arguments...)
}

arguments = append([]string{"./cortex"}, arguments...)

testMode = true
os.Args = arguments
co := captureOutput(t)

// reset default flags
flag.CommandLine = flag.NewFlagSet(arguments[0], flag.ExitOnError)

main()

stdout, stderr := co.Done()
if !bytes.Contains(stdout, stdoutMessage) {
t.Errorf("Expected on stdout: %q, stdout: %s\n", stdoutMessage, stdout)
}
if !bytes.Contains(stderr, stderrMessage) {
t.Errorf("Expected on stderr: %q, stderr: %s\n", stderrMessage, stderr)
}
}

type capturedOutput struct {
stdoutBuf bytes.Buffer
stderrBuf bytes.Buffer

wg sync.WaitGroup
stdoutReader, stdoutWriter *os.File
stderrReader, stderrWriter *os.File
}

func captureOutput(t *testing.T) *capturedOutput {
stdoutR, stdoutW, err := os.Pipe()
require.NoError(t, err)
os.Stdout = stdoutW

stderrR, stderrW, err := os.Pipe()
require.NoError(t, err)
os.Stderr = stderrW

co := &capturedOutput{
stdoutReader: stdoutR,
stdoutWriter: stdoutW,
stderrReader: stderrR,
stderrWriter: stderrW,
}
co.wg.Add(1)
go func() {
defer co.wg.Done()
_, _ = io.Copy(&co.stdoutBuf, stdoutR)
}()

co.wg.Add(1)
go func() {
defer co.wg.Done()
_, _ = io.Copy(&co.stderrBuf, stderrR)
}()

return co
}

func (co *capturedOutput) Done() (stdout []byte, stderr []byte) {
// we need to close writers for readers to stop
_ = co.stdoutWriter.Close()
_ = co.stderrWriter.Close()

co.wg.Wait()

return co.stdoutBuf.Bytes(), co.stderrBuf.Bytes()
}
4 changes: 4 additions & 0 deletions pkg/cortex/modules.go
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,10 @@ func (m *moduleName) Set(s string) error {
}
}

func (m moduleName) MarshalYAML() (interface{}, error) {
return m.String(), nil
}

func (m *moduleName) UnmarshalYAML(unmarshal func(interface{}) error) error {
var s string
if err := unmarshal(&s); err != nil {
Expand Down
2 changes: 1 addition & 1 deletion pkg/ring/kv/multi.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ type MultiConfig struct {
MirrorTimeout time.Duration `yaml:"mirror_timeout"`

// ConfigProvider returns channel with MultiRuntimeConfig updates.
ConfigProvider func() <-chan MultiRuntimeConfig
ConfigProvider func() <-chan MultiRuntimeConfig `yaml:"-"`
}

// RegisterFlagsWithPrefix registers flags with prefix.
Expand Down
22 changes: 22 additions & 0 deletions pkg/util/flagext/ignored.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
package flagext

import (
"flag"
)

type ignoredFlag struct {
name string
}

func (ignoredFlag) String() string {
return "ignored"
}

func (d ignoredFlag) Set(string) error {
return nil
}

// IgnoredFlag ignores set value, without any warning
func IgnoredFlag(f *flag.FlagSet, name, message string) {
f.Var(ignoredFlag{name}, name, message)
}
2 changes: 1 addition & 1 deletion pkg/util/runtimeconfig/manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ type Loader func(filename string) (interface{}, error)
type ManagerConfig struct {
ReloadPeriod time.Duration `yaml:"period"`
LoadPath string `yaml:"file"`
Loader Loader
Loader Loader `yaml:"-"`
}

// RegisterFlags registers flags.
Expand Down