Skip to content
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