Avoid calling flag.Parse() twice. (#1997)

* Calling flag.Parse() twice confuses some flags.

flag.Parse() was called twice previously, first to get -config.file
value to parse the config file, and second to parse remaining command
line parameters – overwriting values from config file.

Unfortunately that confuses command line parameters that accept multiple
values. Here we get -config.file via a separate FlagSet to avoid calling
flag.Parse() (on default FlagSet) twice.

Signed-off-by: Peter Štibraný <[email protected]>

* Reverted unintended import order change.

Signed-off-by: Peter Štibraný <[email protected]>

* Call RegisterFlags(&cfg) to set defaults before parsing config file.

Signed-off-by: Peter Štibraný <[email protected]>

* Added some unit tests for parsing config.

Signed-off-by: Peter Štibraný <[email protected]>

* Test mode in main simply dumps the config YAML if no error occurs

Added tests for happy path.

Signed-off-by: Peter Štibraný <[email protected]>

* Separate stdout/stderr expected messages and checks.

Signed-off-by: Peter Štibraný <[email protected]>

* Use mutexProfileFraction only after flags are parsed.

Signed-off-by: Peter Štibraný <[email protected]>

* Renamed default test name.

Signed-off-by: Peter Štibraný <[email protected]>

* Reset testMode to original value at the end of test.

Signed-off-by: Peter Štibraný <[email protected]>

Author: pstibrany

cmd/cortex/main.go

@@ -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()
@@ -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)
@@ -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)
+	}
+}

cmd/cortex/main_test.go

@@ -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()
+}

pkg/cortex/modules.go

@@ -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 {

pkg/ring/kv/multi.go

@@ -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.

pkg/util/flagext/ignored.go

@@ -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)
+}

pkg/util/runtimeconfig/manager.go

@@ -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.