357 lines
9.1 KiB
Go
357 lines
9.1 KiB
Go
// Copyright (C) 2019 Storj Labs, Inc.
|
|
// See LICENSE for copying information.
|
|
|
|
package process
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"flag"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"os/signal"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/spf13/cobra"
|
|
"github.com/spf13/pflag"
|
|
"github.com/spf13/viper"
|
|
"github.com/zeebo/errs"
|
|
"github.com/zeebo/structs"
|
|
"go.uber.org/zap"
|
|
monkit "gopkg.in/spacemonkeygo/monkit.v2"
|
|
"gopkg.in/spacemonkeygo/monkit.v2/collect"
|
|
"gopkg.in/spacemonkeygo/monkit.v2/present"
|
|
|
|
"storj.io/storj/pkg/cfgstruct"
|
|
"storj.io/storj/private/version"
|
|
)
|
|
|
|
// DefaultCfgFilename is the default filename used for storing a configuration.
|
|
const DefaultCfgFilename = "config.yaml"
|
|
|
|
var (
|
|
mon = monkit.Package()
|
|
|
|
commandMtx sync.Mutex
|
|
contexts = map[*cobra.Command]context.Context{}
|
|
cancels = map[*cobra.Command]context.CancelFunc{}
|
|
configs = map[*cobra.Command][]interface{}{}
|
|
vipers = map[*cobra.Command]*viper.Viper{}
|
|
)
|
|
|
|
// Bind sets flags on a command that match the configuration struct
|
|
// 'config'. It ensures that the config has all of the values loaded into it
|
|
// when the command runs.
|
|
func Bind(cmd *cobra.Command, config interface{}, opts ...cfgstruct.BindOpt) {
|
|
commandMtx.Lock()
|
|
defer commandMtx.Unlock()
|
|
|
|
cfgstruct.Bind(cmd.Flags(), config, opts...)
|
|
configs[cmd] = append(configs[cmd], config)
|
|
}
|
|
|
|
// Exec runs a Cobra command. If a "config-dir" flag is defined it will be parsed
|
|
// and loaded using viper.
|
|
func Exec(cmd *cobra.Command) {
|
|
ExecWithCustomConfig(cmd, true, LoadConfig)
|
|
}
|
|
|
|
// ExecCustomDebug runs default configuration except the default debug is disabled.
|
|
func ExecCustomDebug(cmd *cobra.Command) {
|
|
ExecWithCustomConfig(cmd, false, LoadConfig)
|
|
}
|
|
|
|
// ExecWithCustomConfig runs a Cobra command. Custom configuration can be loaded.
|
|
func ExecWithCustomConfig(cmd *cobra.Command, debugEnabled bool, loadConfig func(cmd *cobra.Command, vip *viper.Viper) error) {
|
|
cmd.AddCommand(&cobra.Command{
|
|
Use: "version",
|
|
Short: "output the version's build information, if any",
|
|
RunE: cmdVersion,
|
|
Annotations: map[string]string{"type": "setup"}})
|
|
|
|
exe, err := os.Executable()
|
|
if err == nil {
|
|
cmd.Use = exe
|
|
}
|
|
|
|
pflag.CommandLine.AddGoFlagSet(flag.CommandLine)
|
|
cleanup(cmd, debugEnabled, loadConfig)
|
|
err = cmd.Execute()
|
|
if err != nil {
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
// Ctx returns the appropriate context.Context for ExecuteWithConfig commands
|
|
func Ctx(cmd *cobra.Command) (context.Context, context.CancelFunc) {
|
|
commandMtx.Lock()
|
|
defer commandMtx.Unlock()
|
|
|
|
ctx := contexts[cmd]
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
contexts[cmd] = ctx
|
|
}
|
|
|
|
cancel := cancels[cmd]
|
|
if cancel == nil {
|
|
ctx, cancel = context.WithCancel(ctx)
|
|
contexts[cmd] = ctx
|
|
cancels[cmd] = cancel
|
|
|
|
c := make(chan os.Signal, 1)
|
|
signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
|
|
go func() {
|
|
sig := <-c
|
|
log.Printf("Got a signal from the OS: %q", sig)
|
|
signal.Stop(c)
|
|
cancel()
|
|
}()
|
|
}
|
|
|
|
return ctx, cancel
|
|
}
|
|
|
|
// Viper returns the appropriate *viper.Viper for the command, creating if necessary.
|
|
func Viper(cmd *cobra.Command) (*viper.Viper, error) {
|
|
return ViperWithCustomConfig(cmd, LoadConfig)
|
|
}
|
|
|
|
// ViperWithCustomConfig returns the appropriate *viper.Viper for the command, creating if necessary. Custom
|
|
// config load logic can be defined with "loadConfig" parameter.
|
|
func ViperWithCustomConfig(cmd *cobra.Command, loadConfig func(cmd *cobra.Command, vip *viper.Viper) error) (*viper.Viper, error) {
|
|
commandMtx.Lock()
|
|
defer commandMtx.Unlock()
|
|
|
|
if vip := vipers[cmd]; vip != nil {
|
|
return vip, nil
|
|
}
|
|
|
|
vip := viper.New()
|
|
if err := vip.BindPFlags(cmd.Flags()); err != nil {
|
|
return nil, err
|
|
}
|
|
vip.SetEnvPrefix("storj")
|
|
vip.SetEnvKeyReplacer(strings.NewReplacer(".", "_", "-", "_"))
|
|
vip.AutomaticEnv()
|
|
|
|
err := loadConfig(cmd, vip)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
vipers[cmd] = vip
|
|
return vip, nil
|
|
}
|
|
|
|
// LoadConfig loads configuration into *viper.Viper from file specified with "config-dir" flag.
|
|
func LoadConfig(cmd *cobra.Command, vip *viper.Viper) error {
|
|
cfgFlag := cmd.Flags().Lookup("config-dir")
|
|
if cfgFlag != nil && cfgFlag.Value.String() != "" {
|
|
path := filepath.Join(os.ExpandEnv(cfgFlag.Value.String()), DefaultCfgFilename)
|
|
if fileExists(path) {
|
|
setupCommand := cmd.Annotations["type"] == "setup"
|
|
vip.SetConfigFile(path)
|
|
if err := vip.ReadInConfig(); err != nil && !setupCommand {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
var traceOut = flag.String("debug.trace-out", "", "If set, a path to write a process trace SVG to")
|
|
|
|
func cleanup(cmd *cobra.Command, debugEnabled bool, loadConfig func(cmd *cobra.Command, vip *viper.Viper) error) {
|
|
for _, ccmd := range cmd.Commands() {
|
|
cleanup(ccmd, debugEnabled, loadConfig)
|
|
}
|
|
if cmd.Run != nil {
|
|
panic("Please use cobra's RunE instead of Run")
|
|
}
|
|
internalRun := cmd.RunE
|
|
if internalRun == nil {
|
|
return
|
|
}
|
|
|
|
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
|
|
ctx := context.Background()
|
|
defer mon.TaskNamed("root")(&ctx)(&err)
|
|
|
|
vip, err := ViperWithCustomConfig(cmd, loadConfig)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
commandMtx.Lock()
|
|
configValues := configs[cmd]
|
|
commandMtx.Unlock()
|
|
|
|
var (
|
|
brokenKeys = map[string]struct{}{}
|
|
missingKeys = map[string]struct{}{}
|
|
usedKeys = map[string]struct{}{}
|
|
allKeys = map[string]struct{}{}
|
|
allSettings = vip.AllSettings()
|
|
)
|
|
|
|
// Hacky hack: these two keys are noprefix which breaks all scoping
|
|
if val, ok := allSettings["api-key"]; ok {
|
|
allSettings["legacy.client.api-key"] = val
|
|
delete(allSettings, "api-key")
|
|
}
|
|
if val, ok := allSettings["satellite-addr"]; ok {
|
|
allSettings["legacy.client.satellite-addr"] = val
|
|
delete(allSettings, "satellite-addr")
|
|
}
|
|
|
|
for _, config := range configValues {
|
|
// Decode and all of the resulting keys into our sets
|
|
res := structs.Decode(allSettings, config)
|
|
for key := range res.Used {
|
|
usedKeys[key] = struct{}{}
|
|
allKeys[key] = struct{}{}
|
|
}
|
|
for key := range res.Missing {
|
|
missingKeys[key] = struct{}{}
|
|
allKeys[key] = struct{}{}
|
|
}
|
|
for key := range res.Broken {
|
|
brokenKeys[key] = struct{}{}
|
|
allKeys[key] = struct{}{}
|
|
}
|
|
}
|
|
|
|
// Propagate keys that are missing to flags, and remove any used keys
|
|
// from the missing set.
|
|
for key := range missingKeys {
|
|
if f := cmd.Flags().Lookup(key); f != nil {
|
|
val := vip.GetString(key)
|
|
err := f.Value.Set(val)
|
|
f.Changed = val != f.DefValue
|
|
if err != nil {
|
|
brokenKeys[key] = struct{}{}
|
|
} else {
|
|
usedKeys[key] = struct{}{}
|
|
}
|
|
} else if f := flag.Lookup(key); f != nil {
|
|
err := f.Value.Set(vip.GetString(key))
|
|
if err != nil {
|
|
brokenKeys[key] = struct{}{}
|
|
} else {
|
|
usedKeys[key] = struct{}{}
|
|
}
|
|
}
|
|
}
|
|
for key := range missingKeys {
|
|
if _, ok := usedKeys[key]; ok {
|
|
delete(missingKeys, key)
|
|
}
|
|
}
|
|
|
|
logger, err := NewLogger()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if vip.ConfigFileUsed() != "" {
|
|
path, err := filepath.Abs(vip.ConfigFileUsed())
|
|
if err != nil {
|
|
path = vip.ConfigFileUsed()
|
|
logger.Debug("unable to resolve path", zap.Error(err))
|
|
}
|
|
|
|
logger.Sugar().Info("Configuration loaded from: ", path)
|
|
}
|
|
|
|
defer func() { _ = logger.Sync() }()
|
|
defer zap.ReplaceGlobals(logger)()
|
|
defer zap.RedirectStdLog(logger)()
|
|
|
|
// okay now that logging is working, inform about the broken keys
|
|
if cmd.Annotations["type"] != "helper" {
|
|
for key := range missingKeys {
|
|
logger.Sugar().Infof("Invalid configuration file key: %s", key)
|
|
}
|
|
}
|
|
for key := range brokenKeys {
|
|
logger.Sugar().Infof("Invalid configuration file value for key: %s", key)
|
|
}
|
|
|
|
if debugEnabled {
|
|
err = initDebug(logger, monkit.Default)
|
|
if err != nil {
|
|
withoutStack := errors.New(err.Error())
|
|
logger.Debug("failed to start debug endpoints", zap.Error(withoutStack))
|
|
err = nil
|
|
}
|
|
}
|
|
|
|
var workErr error
|
|
work := func(ctx context.Context) {
|
|
commandMtx.Lock()
|
|
contexts[cmd] = ctx
|
|
commandMtx.Unlock()
|
|
defer func() {
|
|
commandMtx.Lock()
|
|
delete(contexts, cmd)
|
|
delete(cancels, cmd)
|
|
commandMtx.Unlock()
|
|
}()
|
|
|
|
workErr = internalRun(cmd, args)
|
|
}
|
|
|
|
if *traceOut != "" {
|
|
fh, err := os.Create(*traceOut)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if strings.HasSuffix(*traceOut, ".json") {
|
|
err = present.SpansToJSON(fh, collect.CollectSpans(ctx, work))
|
|
} else {
|
|
err = present.SpansToSVG(fh, collect.CollectSpans(ctx, work))
|
|
}
|
|
err = errs.Combine(err, fh.Close())
|
|
if err != nil {
|
|
logger.Error("failed to write svg", zap.Error(err))
|
|
}
|
|
} else {
|
|
work(ctx)
|
|
}
|
|
|
|
err = workErr
|
|
if err != nil {
|
|
logger.Debug("Unrecoverable error", zap.Error(err))
|
|
fmt.Println("Error:", err.Error())
|
|
os.Exit(1)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func cmdVersion(cmd *cobra.Command, args []string) (err error) {
|
|
if version.Build.Release {
|
|
fmt.Println("Release build")
|
|
} else {
|
|
fmt.Println("Development build")
|
|
}
|
|
|
|
if !version.Build.Version.IsZero() {
|
|
fmt.Println("Version:", version.Build.Version.String())
|
|
}
|
|
if !version.Build.Timestamp.IsZero() {
|
|
fmt.Println("Build timestamp:", version.Build.Timestamp.Format(time.RFC822))
|
|
}
|
|
if version.Build.CommitHash != "" {
|
|
fmt.Println("Git commit:", version.Build.CommitHash)
|
|
}
|
|
return err
|
|
}
|