Add -modules support (#497)

This adds the ability to list modules, display a basic flag `-help` and `-help-all` including all flags. And handles `-version` from the main package.

Author: 0xPoe

cmd/phlare/main.go

@@ -5,25 +5,96 @@ import (
 	"flag"
 	"fmt"
 	"os"
+	"sort"
+
+	"github.com/grafana/dskit/flagext"
+	"github.com/prometheus/common/version"
 
 	"github.com/grafana/phlare/pkg/cfg"
 	"github.com/grafana/phlare/pkg/phlare"
+	"github.com/grafana/phlare/pkg/usage"
 	_ "github.com/grafana/phlare/pkg/util/build"
 )
 
+type mainFlags struct {
+	phlare.Config
+
+	PrintVersion bool
+	PrintModules bool
+	PrintHelp    bool
+	PrintHelpAll bool
+}
+
+func (mf *mainFlags) Clone() flagext.Registerer {
+	return func(mf mainFlags) *mainFlags {
+		return &mf
+	}(*mf)
+}
+
+func (mf *mainFlags) PhlareConfig() *phlare.Config {
+	return &mf.Config
+}
+
+func (mf *mainFlags) RegisterFlags(fs *flag.FlagSet) {
+	mf.Config.RegisterFlags(fs)
+	fs.BoolVar(&mf.PrintVersion, "version", false, "Show the version of phlare and exit")
+	fs.BoolVar(&mf.PrintModules, "modules", false, "List available modules that can be used as target and exit.")
+	fs.BoolVar(&mf.PrintHelp, "h", false, "Print basic help.")
+	fs.BoolVar(&mf.PrintHelp, "help", false, "Print basic help.")
+	fs.BoolVar(&mf.PrintHelpAll, "help-all", false, "Print help, also including advanced and experimental parameters.")
+}
+
 func main() {
-	var config phlare.Config
-	if err := cfg.DynamicUnmarshal(&config, os.Args[1:], flag.CommandLine); err != nil {
+	var (
+		flags mainFlags
+	)
+
+	if err := cfg.DynamicUnmarshal(&flags, os.Args[1:], flag.CommandLine); err != nil {
 		fmt.Fprintf(os.Stderr, "failed parsing config: %v\n", err)
 		os.Exit(1)
 	}
 
-	f, err := phlare.New(config)
+	f, err := phlare.New(flags.Config)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "failed creating phlare: %v\n", err)
 		os.Exit(1)
 	}
 
+	if flags.PrintVersion {
+		fmt.Println(version.Print("phlare"))
+		return
+	}
+
+	if flags.PrintModules {
+		allDeps := f.ModuleManager.DependenciesForModule(phlare.All)
+
+		for _, m := range f.ModuleManager.UserVisibleModuleNames() {
+			ix := sort.SearchStrings(allDeps, m)
+			included := ix < len(allDeps) && allDeps[ix] == m
+
+			if included {
+				fmt.Fprintln(os.Stdout, m, "*")
+			} else {
+				fmt.Fprintln(os.Stdout, m)
+			}
+		}
+
+		fmt.Fprintln(os.Stdout)
+		fmt.Fprintln(os.Stdout, "Modules marked with * are included in target All.")
+		return
+	}
+
+	if flags.PrintHelp || flags.PrintHelpAll {
+		// Print available parameters to stdout, so that users can grep/less them easily.
+		flag.CommandLine.SetOutput(os.Stdout)
+		if err := usage.Usage(flags.PrintHelpAll, &flags); err != nil {
+			fmt.Fprintf(os.Stderr, "error printing usage: %s\n", err)
+			os.Exit(1)
+		}
+
+		return
+	}
+
 	err = f.Run()
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "failed running phlare: %v\n", err)

cmd/phlare/main_test.go

@@ -0,0 +1,103 @@
+package main
+
+import (
+	"flag"
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus"
+
+	"github.com/grafana/phlare/pkg/test"
+)
+
+func TestFlagParsing(t *testing.T) {
+	for name, tc := range map[string]struct {
+		arguments      []string
+		stdoutMessage  string // string that must be included in stdout
+		stderrMessage  string // string that must be included in stderr
+		stdoutExcluded string // string that must NOT be included in stdout
+		stderrExcluded string // string that must NOT be included in stderr
+	}{
+		"help-short": {
+			arguments:      []string{"-h"},
+			stdoutMessage:  "Usage of", // Usage must be on stdout, not stderr.
+			stderrExcluded: "Usage of",
+		},
+		"help": {
+			arguments:      []string{"-help"},
+			stdoutMessage:  "Usage of", // Usage must be on stdout, not stderr.
+			stderrExcluded: "Usage of",
+		},
+		"help-all": {
+			arguments:      []string{"-help-all"},
+			stdoutMessage:  "Usage of", // Usage must be on stdout, not stderr.
+			stderrExcluded: "Usage of",
+		},
+		"user visible module listing": {
+			arguments:      []string{"-modules"},
+			stdoutMessage:  "ingester *\n",
+			stderrExcluded: "ingester\n",
+		},
+		"version": {
+			arguments:      []string{"-version"},
+			stdoutMessage:  "phlare, version",
+			stderrExcluded: "phlare, version",
+		},
+	} {
+		t.Run(name, func(t *testing.T) {
+			_ = os.Setenv("TARGET", "ingester")
+			oldDefaultRegistry := prometheus.DefaultRegisterer
+			defer func() {
+				prometheus.DefaultRegisterer = oldDefaultRegistry
+			}()
+			// We need to reset the default registry to avoid
+			// "duplicate metrics collector registration attempted" errors.
+			prometheus.DefaultRegisterer = prometheus.NewRegistry()
+			testSingle(t, tc.arguments, tc.stdoutMessage, tc.stderrMessage, tc.stdoutExcluded, tc.stderrExcluded)
+		})
+	}
+}
+
+func testSingle(t *testing.T, arguments []string, stdoutMessage, stderrMessage, stdoutExcluded, stderrExcluded string) {
+	t.Helper()
+	oldArgs, oldStdout, oldStderr := os.Args, os.Stdout, os.Stderr
+	restored := false
+	restoreIfNeeded := func() {
+		if restored {
+			return
+		}
+		os.Stdout = oldStdout
+		os.Stderr = oldStderr
+		os.Args = oldArgs
+		restored = true
+	}
+	defer restoreIfNeeded()
+
+	arguments = append([]string{"./phlare"}, arguments...)
+
+	os.Args = arguments
+	co := test.CaptureOutput(t)
+
+	// reset default flags
+	flag.CommandLine = flag.NewFlagSet(arguments[0], flag.ExitOnError)
+
+	main()
+
+	stdout, stderr := co.Done()
+
+	// Restore stdout and stderr before reporting errors to make them visible.
+	restoreIfNeeded()
+	if !strings.Contains(stdout, stdoutMessage) {
+		t.Errorf("Expected on stdout: %q, stdout: %s\n", stdoutMessage, stdout)
+	}
+	if !strings.Contains(stderr, stderrMessage) {
+		t.Errorf("Expected on stderr: %q, stderr: %s\n", stderrMessage, stderr)
+	}
+	if len(stdoutExcluded) > 0 && strings.Contains(stdout, stdoutExcluded) {
+		t.Errorf("Unexpected output on stdout: %q, stdout: %s\n", stdoutExcluded, stdout)
+	}
+	if len(stderrExcluded) > 0 && strings.Contains(stderr, stderrExcluded) {
+		t.Errorf("Unexpected output on stderr: %q, stderr: %s\n", stderrExcluded, stderr)
+	}
+}

pkg/cfg/files.go

@@ -74,17 +74,10 @@ func YAMLFlag(args []string, name string) Source {
 		// parsing out the config file location.
 		dst.Clone().RegisterFlags(freshFlags)
 
-		usage := freshFlags.Usage
 		freshFlags.Usage = func() { /* don't do anything by default, we will print usage ourselves, but only when requested. */ }
 
-		err := freshFlags.Parse(args)
-		if err == flag.ErrHelp {
-			// print available parameters to stdout, so that users can grep/less it easily
-			freshFlags.SetOutput(os.Stdout)
-			usage()
-			os.Exit(2)
-		} else if err != nil {
-			fmt.Fprintln(freshFlags.Output(), "Run with -help to get list of available parameters")
+		if err := freshFlags.Parse(args); err != nil {
+			fmt.Fprintln(freshFlags.Output(), "Run with -help to get a list of available parameters")
 			os.Exit(2)
 		}
 
@@ -99,6 +92,5 @@ func YAMLFlag(args []string, name string) Source {
 		}
 
 		return YAML(f.Value.String(), expandEnv)(dst)
-
 	}
 }

pkg/phlare/phlare.go

@@ -75,7 +75,6 @@ type Config struct {
 	Analytics           usagestats.Config `yaml:"analytics"`
 
 	ConfigFile      string `yaml:"-"`
-	ShowVersion     bool   `yaml:"-"`
 	ConfigExpandEnv bool   `yaml:"-"`
 }
 
@@ -106,7 +105,6 @@ func (c *Config) RegisterFlagsWithContext(ctx context.Context, f *flag.FlagSet)
 	f.Var(&c.Target, "target", "Comma-separated list of Phlare modules to load. "+
 		"The alias 'all' can be used in the list to load a number of core modules and will enable single-binary mode. ")
 	f.BoolVar(&c.MultitenancyEnabled, "auth.multitenancy-enabled", false, "When set to true, incoming HTTP requests must specify tenant ID in HTTP X-Scope-OrgId header. When set to false, tenant ID anonymous is used instead.")
-	f.BoolVar(&c.ShowVersion, "version", false, "Show the version of phlare and exit")
 	f.BoolVar(&c.ConfigExpandEnv, "config.expand-env", false, "Expands ${var} in config according to the values of the environment variables.")
 
 	c.registerServerFlagsWithChangedDefaultValues(f)
@@ -165,6 +163,10 @@ func (c *Config) Validate() error {
 	return c.AgentConfig.Validate()
 }
 
+type phlareConfigGetter interface {
+	PhlareConfig() *Config
+}
+
 func (c *Config) ApplyDynamicConfig() cfg.Source {
 	c.Ingester.LifecyclerConfig.RingConfig.KVStore.Store = "memberlist"
 	c.Distributor.DistributorRing.KVStore.Store = c.Ingester.LifecyclerConfig.RingConfig.KVStore.Store
@@ -175,10 +177,11 @@ func (c *Config) ApplyDynamicConfig() cfg.Source {
 	c.Worker.MaxConcurrentRequests = 4 // todo we might want this as a config flags.
 
 	return func(dst cfg.Cloneable) error {
-		r, ok := dst.(*Config)
+		g, ok := dst.(phlareConfigGetter)
 		if !ok {
-			return errors.New("dst is not a Phlare config")
+			return fmt.Errorf("dst is not a Phlare config getter %T", dst)
 		}
+		r := g.PhlareConfig()
 		if r.AgentConfig.ClientConfig.URL.String() == "" {
 			listenAddress := "0.0.0.0"
 			if c.Server.HTTPListenAddress != "" {
@@ -233,11 +236,6 @@ func New(cfg Config) (*Phlare, error) {
 	logger := initLogger(&cfg.Server)
 	usagestats.Edition("oss")
 
-	if cfg.ShowVersion {
-		fmt.Println(version.Print("phlare"))
-		os.Exit(0)
-	}
-
 	phlare := &Phlare{
 		Cfg:    cfg,
 		logger: logger,

pkg/phlare/phlare_test.go

@@ -22,6 +22,9 @@ func TestFlagDefaults(t *testing.T) {
 	f.PrintDefaults()
 
 	const delim = '\n'
+	// Because this is a short flag, it will be printed on the same line as the
+	// flag name. So we need to ignore this special case.
+	const ignoredHelpFlags = "-h\tPrint basic help."
 
 	// Populate map with parsed default flags.
 	// Key is the flag and value is the default text.
@@ -33,6 +36,10 @@ func TestFlagDefaults(t *testing.T) {
 		}
 		require.NoError(t, err)
 
+		if strings.Contains(line, ignoredHelpFlags) {
+			continue
+		}
+
 		nextLine, err := buf.ReadString(delim)
 		require.NoError(t, err)
 

pkg/test/capture.go

@@ -0,0 +1,68 @@
+package test
+
+import (
+	"bytes"
+	"io"
+	"os"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type CapturedOutput struct {
+	stdoutBuf bytes.Buffer
+	stderrBuf bytes.Buffer
+
+	wg                         sync.WaitGroup
+	stdoutReader, stdoutWriter *os.File
+	stderrReader, stderrWriter *os.File
+}
+
+// CaptureOutput replaces os.Stdout and os.Stderr with new pipes, that will
+// write output to buffers. Buffers are accessible by calling Done on returned
+// struct.
+//
+// os.Stdout and os.Stderr must be reverted to previous values manually.
+func CaptureOutput(t *testing.T) *CapturedOutput {
+	stdoutR, stdoutW, err := os.Pipe()
+	require.NoError(t, err)
+
+	stderrR, stderrW, err := os.Pipe()
+	require.NoError(t, err)
+
+	os.Stdout = stdoutW
+	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
+}
+
+// Done waits until all captured output has been written to buffers,
+// and then returns the buffers.
+func (co *CapturedOutput) Done() (stdout string, stderr string) {
+	// we need to close writers for readers to stop
+	_ = co.stdoutWriter.Close()
+	_ = co.stderrWriter.Close()
+
+	co.wg.Wait()
+
+	return co.stdoutBuf.String(), co.stderrBuf.String()
+}