LBRYCommunity/tracker
1
0
Fork 0
Code Issues 1 Pull requests Projects Releases Packages Wiki Activity

cmd/chihaya: pull out pre and post funcs

Browse source
This commit is contained in:
Jimmy Zelinskie 2017-12-22 23:53:58 -05:00
parent 5840cd3de1
commit 2bead6b7b4
1 changed files with 60 additions and 50 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

110
cmd/chihaya/main.go
View file

@ -167,59 +167,69 @@ func RunCmdFunc(cmd *cobra.Command, args []string) error {
} }
} }
// PreRunCmdFunc handles command line flags for the Run command.
func PreRunCmdFunc(cmd *cobra.Command, args []string) error {
noColors, err := cmd.Flags().GetBool("nocolors")
if err != nil {
return err
}
if noColors {
log.SetFormatter(&logrus.TextFormatter{DisableColors: true})
}
jsonLog, err := cmd.Flags().GetBool("json")
if err != nil {
return err
}
if jsonLog {
log.SetFormatter(&logrus.JSONFormatter{})
log.Info("enabled JSON logging")
}
debugLog, err := cmd.Flags().GetBool("debug")
if err != nil {
return err
}
if debugLog {
log.SetDebug(true)
log.Info("enabled debug logging")
}
cpuProfilePath, err := cmd.Flags().GetString("cpuprofile")
if err != nil {
return err
}
if cpuProfilePath != "" {
f, err := os.Create(cpuProfilePath)
if err != nil {
return err
}
pprof.StartCPUProfile(f)
log.Info("enabled CPU profiling", log.Fields{"path": cpuProfilePath})
}
return nil
}
// PostRunCmdFunc handles clean up of any state initialized by command line
// flags.
func PostRunCmdFunc(cmd *cobra.Command, args []string) error {
// This can be called regardless because it noops when not profiling.
pprof.StopCPUProfile()
return nil
}
func main() { func main() {
var rootCmd = &cobra.Command{ var rootCmd = &cobra.Command{
Use: "chihaya", Use: "chihaya",
Short: "BitTorrent Tracker", Short: "BitTorrent Tracker",
Long: "A customizable, multi-protocol BitTorrent Tracker", Long: "A customizable, multi-protocol BitTorrent Tracker",
PersistentPreRunE: func(cmd *cobra.Command, args []string) error { PersistentPreRunE: PreRunCmdFunc,
noColors, err := cmd.Flags().GetBool("nocolors") RunE: RunCmdFunc,
if err != nil { PersistentPostRunE: PostRunCmdFunc,
return err
}
if noColors {
log.SetFormatter(&logrus.TextFormatter{DisableColors: true})
}
jsonLog, err := cmd.Flags().GetBool("json")
if err != nil {
return err
}
if jsonLog {
log.SetFormatter(&logrus.JSONFormatter{})
}
debugLog, err := cmd.Flags().GetBool("debug")
if err != nil {
return err
}
if debugLog {
log.Info("enabling debug logging")
log.SetDebug(true)
}
cpuProfilePath, err := cmd.Flags().GetString("cpuprofile")
if err != nil {
return err
}
if cpuProfilePath != "" {
log.Info("enabling CPU profiling", log.Fields{"path": cpuProfilePath})
f, err := os.Create(cpuProfilePath)
if err != nil {
return err
}
pprof.StartCPUProfile(f)
}
return nil
},
RunE: RunCmdFunc,
PersistentPostRunE: func(cmd *cobra.Command, args []string) error {
// StopCPUProfile() noops when not profiling.
pprof.StopCPUProfile()
return nil
},
} }
rootCmd.Flags().String("config", "/etc/chihaya.yaml", "location of configuration file") rootCmd.Flags().String("config", "/etc/chihaya.yaml", "location of configuration file")
rootCmd.Flags().String("cpuprofile", "", "location to save a CPU profile") rootCmd.Flags().String("cpuprofile", "", "location to save a CPU profile")
rootCmd.Flags().Bool("debug", false, "enable debug logging") rootCmd.Flags().Bool("debug", false, "enable debug logging")