Просмотр исходного кода

Update logging package and flag parsing

Daniel Ramich 4 лет назад
Родитель
Сommit
35cefbaa2c
3 измененных файлов с 62 добавлено и 25 удалено
  1. 2 5
      cmd/costmodel/main.go
  2. 17 11
      pkg/cmd/commands.go
  3. 43 9
      pkg/log/log.go

+ 2 - 5
cmd/costmodel/main.go

@@ -1,17 +1,14 @@
 package main
 
 import (
-	"os"
-
 	"github.com/kubecost/cost-model/pkg/cmd"
-	"k8s.io/klog"
+	"github.com/rs/zerolog/log"
 )
 
 func main() {
 	// runs the appropriate application mode using the default cost-model command
 	// see: github.com/kubecost/cost-model/pkg/cmd package for details
 	if err := cmd.Execute(nil); err != nil {
-		klog.Fatal(err)
-		os.Exit(1)
+		log.Fatal().Err(err)
 	}
 }

+ 17 - 11
pkg/cmd/commands.go

@@ -1,15 +1,15 @@
 package cmd
 
 import (
-	"flag"
 	"fmt"
 	"os"
+	"strings"
 
 	"github.com/kubecost/cost-model/pkg/cmd/agent"
 	"github.com/kubecost/cost-model/pkg/cmd/costmodel"
+	"github.com/kubecost/cost-model/pkg/log"
 	"github.com/spf13/cobra"
-	"github.com/spf13/pflag"
-	"k8s.io/klog"
+	"github.com/spf13/viper"
 )
 
 const (
@@ -41,14 +41,6 @@ func Execute(costModelCmd *cobra.Command) error {
 
 	rootCmd := newRootCommand(costModelCmd)
 
-	// initialize klog and make cobra aware of all the go flags
-	klog.InitFlags(nil)
-
-	flag.CommandLine.VisitAll(func(f *flag.Flag) {
-		pflag.CommandLine.AddGoFlag(f)
-	})
-	pflag.CommandLine.Set("v", "3")
-
 	// in the event that no directive/command is passed, we want to default to using the cost-model command
 	// cobra doesn't provide a way within the API to do this, so we'll prepend the command if it is omitted.
 	if len(os.Args) > 1 {
@@ -73,12 +65,26 @@ func newRootCommand(costModelCmd *cobra.Command) *cobra.Command {
 		SilenceUsage: true,
 	}
 
+	// Add our persistent flags, these are global and available anywhere
+	cmd.PersistentFlags().String("log-level", "info", "Set the log level")
+	cmd.PersistentFlags().String("log-format", "pretty", "Set the log format - Can be either 'JSON' or 'pretty'")
+
+	viper.BindPFlag("log-level", cmd.PersistentFlags().Lookup("log-level"))
+	viper.BindPFlag("log-format", cmd.PersistentFlags().Lookup("log-format"))
+
+	// Setup viper to read from the env, this allows reading flags from the command line or the env
+	// using the format 'LOG_LEVEL'
+	viper.AutomaticEnv()
+	viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
+
 	// add the modes of operation
 	cmd.AddCommand(
 		costModelCmd,
 		newAgentCommand(),
 	)
 
+	log.InitLogging()
+
 	return cmd
 }
 

+ 43 - 9
pkg/log/log.go

@@ -2,9 +2,13 @@ package log
 
 import (
 	"fmt"
+	"os"
+	"strings"
 	"time"
 
-	"k8s.io/klog"
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
+	"github.com/spf13/viper"
 )
 
 // TODO for deduped functions, if timeLogged > logTypeLimit, should we log once
@@ -13,8 +17,26 @@ import (
 // concurrency-safe counter
 var ctr = newCounter()
 
+func InitLogging() {
+	zerolog.TimeFieldFormat = time.RFC3339Nano
+	// Default to using pretty formatting
+	if strings.ToLower(viper.GetString("log-format")) != "json" {
+		log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr, TimeFormat: time.RFC3339Nano})
+	}
+
+	level, err := zerolog.ParseLevel(viper.GetString("log-level"))
+	if err != nil {
+		zerolog.SetGlobalLevel(zerolog.InfoLevel)
+		log.Warn().Msg("Error parsing log-level, setting level to 'info'")
+		return
+	}
+	zerolog.SetGlobalLevel(level)
+	log.Log().Msgf("Log level set to %v", level)
+
+}
+
 func Errorf(format string, a ...interface{}) {
-	klog.Errorf(fmt.Sprintf("[Error] %s", format), a...)
+	log.Error().Msgf(format, a...)
 }
 
 func DedupedErrorf(logTypeLimit int, format string, a ...interface{}) {
@@ -28,23 +50,27 @@ func DedupedErrorf(logTypeLimit int, format string, a ...interface{}) {
 	}
 }
 
-func Warningf(format string, a ...interface{}) {
-	klog.V(2).Infof(fmt.Sprintf("[Warning] %s", format), a...)
+func Warnf(format string, a ...interface{}) {
+	log.Warn().Msgf(format, a...)
 }
 
 func DedupedWarningf(logTypeLimit int, format string, a ...interface{}) {
 	timesLogged := ctr.increment(format)
 
 	if timesLogged < logTypeLimit {
-		Warningf(format, a...)
+		Warnf(format, a...)
 	} else if timesLogged == logTypeLimit {
-		Warningf(format, a...)
+		Warnf(format, a...)
 		Infof("%s logged %d times: suppressing future logs", format, logTypeLimit)
 	}
 }
 
+func Info(msg string) {
+	log.Info().Msg(msg)
+}
+
 func Infof(format string, a ...interface{}) {
-	klog.V(3).Infof(fmt.Sprintf("[Info] %s", format), a...)
+	log.Info().Msgf(format, a...)
 }
 
 func DedupedInfof(logTypeLimit int, format string, a ...interface{}) {
@@ -59,11 +85,19 @@ func DedupedInfof(logTypeLimit int, format string, a ...interface{}) {
 }
 
 func Profilef(format string, a ...interface{}) {
-	klog.V(3).Infof(fmt.Sprintf("[Profiler] %s", format), a...)
+	log.Info().Msgf(fmt.Sprintf("[Profiler] %s", format), a...)
+}
+
+func Debug(msg string) {
+	log.Debug().Msg(msg)
 }
 
 func Debugf(format string, a ...interface{}) {
-	klog.V(5).Infof(fmt.Sprintf("[Debug] %s", format), a...)
+	log.Debug().Msgf(format, a...)
+}
+
+func Fatalf(format string, a ...interface{}) {
+	log.Fatal().Msgf(format, a...)
 }
 
 func Profile(start time.Time, name string) {