github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/benchmark/cmd/command/root.go (about)

     1  package command
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"runtime"
     7  	"strings"
     8  
     9  	"github.com/pyroscope-io/pyroscope/benchmark/internal/config"
    10  	"github.com/pyroscope-io/pyroscope/pkg/cli"
    11  	"github.com/sirupsen/logrus"
    12  	"github.com/spf13/cobra"
    13  )
    14  
    15  func newRootCmd(*config.Config) *cobra.Command {
    16  	rootCmd := &cobra.Command{
    17  		Use: "pyrobench [flags] <subcommand>",
    18  	}
    19  
    20  	rootCmd.SetUsageFunc(func(cmd *cobra.Command) error {
    21  		fmt.Println(gradientBanner())
    22  		fmt.Println(cli.DefaultUsageFunc(cmd.Flags(), cmd))
    23  		return nil
    24  	})
    25  
    26  	rootCmd.SetHelpFunc(func(cmd *cobra.Command, a []string) {
    27  		fmt.Println(gradientBanner())
    28  		fmt.Println(cli.DefaultUsageFunc(cmd.Flags(), cmd))
    29  	})
    30  
    31  	return rootCmd
    32  }
    33  
    34  // Initialize adds all child commands to the root command and sets flags appropriately
    35  func Initialize() error {
    36  	var cfg config.Config
    37  
    38  	rootCmd := newRootCmd(&cfg)
    39  	rootCmd.SilenceErrors = true
    40  	rootCmd.AddCommand(
    41  		newLoadGen(&cfg.LoadGen),
    42  		newPromQuery(&cfg.PromQuery),
    43  		newReport(&cfg.Report),
    44  	)
    45  
    46  	logrus.SetReportCaller(true)
    47  	logrus.SetFormatter(&logrus.TextFormatter{
    48  		TimestampFormat: "2006-01-02T15:04:05.000000",
    49  		FullTimestamp:   true,
    50  		CallerPrettyfier: func(f *runtime.Frame) (string, string) {
    51  			filename := f.File
    52  			if len(filename) > 38 {
    53  				filename = filename[38:]
    54  			}
    55  			return "", fmt.Sprintf(" %s:%d", filename, f.Line)
    56  		},
    57  	})
    58  
    59  	args := os.Args[1:]
    60  	for i, arg := range args {
    61  		if len(arg) > 2 && strings.HasPrefix(arg, "-") && !strings.HasPrefix(arg, "--") {
    62  			args[i] = fmt.Sprintf("-%s", arg)
    63  		}
    64  	}
    65  
    66  	rootCmd.SetArgs(args)
    67  	return rootCmd.Execute()
    68  }