github.com/jd-ly/tools@v0.5.7/internal/tool/tool.go (about) 1 // Copyright 2018 The Go Authors. All rights reserved. 2 // Use of this source code is governed by a BSD-style 3 // license that can be found in the LICENSE file. 4 5 // Package tool is a harness for writing Go tools. 6 package tool 7 8 import ( 9 "context" 10 "flag" 11 "fmt" 12 "log" 13 "os" 14 "reflect" 15 "runtime" 16 "runtime/pprof" 17 "runtime/trace" 18 "time" 19 ) 20 21 // This file is a harness for writing your main function. 22 // The original version of the file is in github.com/jd-ly/tools/internal/tool. 23 // 24 // It adds a method to the Application type 25 // Main(name, usage string, args []string) 26 // which should normally be invoked from a true main as follows: 27 // func main() { 28 // (&Application{}).Main("myapp", "non-flag-command-line-arg-help", os.Args[1:]) 29 // } 30 // It recursively scans the application object for fields with a tag containing 31 // `flag:"flagname" help:"short help text"`` 32 // uses all those fields to build command line flags. 33 // It expects the Application type to have a method 34 // Run(context.Context, args...string) error 35 // which it invokes only after all command line flag processing has been finished. 36 // If Run returns an error, the error will be printed to stderr and the 37 // application will quit with a non zero exit status. 38 39 // Profile can be embedded in your application struct to automatically 40 // add command line arguments and handling for the common profiling methods. 41 type Profile struct { 42 CPU string `flag:"profile.cpu" help:"write CPU profile to this file"` 43 Memory string `flag:"profile.mem" help:"write memory profile to this file"` 44 Trace string `flag:"profile.trace" help:"write trace log to this file"` 45 } 46 47 // Application is the interface that must be satisfied by an object passed to Main. 48 type Application interface { 49 // Name returns the application's name. It is used in help and error messages. 50 Name() string 51 // Most of the help usage is automatically generated, this string should only 52 // describe the contents of non flag arguments. 53 Usage() string 54 // ShortHelp returns the one line overview of the command. 55 ShortHelp() string 56 // DetailedHelp should print a detailed help message. It will only ever be shown 57 // when the ShortHelp is also printed, so there is no need to duplicate 58 // anything from there. 59 // It is passed the flag set so it can print the default values of the flags. 60 // It should use the flag sets configured Output to write the help to. 61 DetailedHelp(*flag.FlagSet) 62 // Run is invoked after all flag processing, and inside the profiling and 63 // error handling harness. 64 Run(ctx context.Context, args ...string) error 65 } 66 67 // This is the type returned by CommandLineErrorf, which causes the outer main 68 // to trigger printing of the command line help. 69 type commandLineError string 70 71 func (e commandLineError) Error() string { return string(e) } 72 73 // CommandLineErrorf is like fmt.Errorf except that it returns a value that 74 // triggers printing of the command line help. 75 // In general you should use this when generating command line validation errors. 76 func CommandLineErrorf(message string, args ...interface{}) error { 77 return commandLineError(fmt.Sprintf(message, args...)) 78 } 79 80 // Main should be invoked directly by main function. 81 // It will only return if there was no error. If an error 82 // was encountered it is printed to standard error and the 83 // application exits with an exit code of 2. 84 func Main(ctx context.Context, app Application, args []string) { 85 s := flag.NewFlagSet(app.Name(), flag.ExitOnError) 86 s.Usage = func() { 87 fmt.Fprint(s.Output(), app.ShortHelp()) 88 fmt.Fprintf(s.Output(), "\n\nUsage: %v [flags] %v\n", app.Name(), app.Usage()) 89 app.DetailedHelp(s) 90 } 91 if err := Run(ctx, app, args); err != nil { 92 fmt.Fprintf(s.Output(), "%s: %v\n", app.Name(), err) 93 if _, printHelp := err.(commandLineError); printHelp { 94 s.Usage() 95 } 96 os.Exit(2) 97 } 98 } 99 100 // Run is the inner loop for Main; invoked by Main, recursively by 101 // Run, and by various tests. It runs the application and returns an 102 // error. 103 func Run(ctx context.Context, app Application, args []string) error { 104 s := flag.NewFlagSet(app.Name(), flag.ExitOnError) 105 s.Usage = func() { 106 fmt.Fprint(s.Output(), app.ShortHelp()) 107 fmt.Fprintf(s.Output(), "\n\nUsage: %v [flags] %v\n", app.Name(), app.Usage()) 108 app.DetailedHelp(s) 109 } 110 p := addFlags(s, reflect.StructField{}, reflect.ValueOf(app)) 111 s.Parse(args) 112 113 if p != nil && p.CPU != "" { 114 f, err := os.Create(p.CPU) 115 if err != nil { 116 return err 117 } 118 if err := pprof.StartCPUProfile(f); err != nil { 119 return err 120 } 121 defer pprof.StopCPUProfile() 122 } 123 124 if p != nil && p.Trace != "" { 125 f, err := os.Create(p.Trace) 126 if err != nil { 127 return err 128 } 129 if err := trace.Start(f); err != nil { 130 return err 131 } 132 defer func() { 133 trace.Stop() 134 log.Printf("To view the trace, run:\n$ go tool trace view %s", p.Trace) 135 }() 136 } 137 138 if p != nil && p.Memory != "" { 139 f, err := os.Create(p.Memory) 140 if err != nil { 141 return err 142 } 143 defer func() { 144 runtime.GC() // get up-to-date statistics 145 if err := pprof.WriteHeapProfile(f); err != nil { 146 log.Printf("Writing memory profile: %v", err) 147 } 148 f.Close() 149 }() 150 } 151 152 return app.Run(ctx, s.Args()...) 153 } 154 155 // addFlags scans fields of structs recursively to find things with flag tags 156 // and add them to the flag set. 157 func addFlags(f *flag.FlagSet, field reflect.StructField, value reflect.Value) *Profile { 158 // is it a field we are allowed to reflect on? 159 if field.PkgPath != "" { 160 return nil 161 } 162 // now see if is actually a flag 163 flagName, isFlag := field.Tag.Lookup("flag") 164 help := field.Tag.Get("help") 165 if !isFlag { 166 // not a flag, but it might be a struct with flags in it 167 if value.Elem().Kind() != reflect.Struct { 168 return nil 169 } 170 p, _ := value.Interface().(*Profile) 171 // go through all the fields of the struct 172 sv := value.Elem() 173 for i := 0; i < sv.Type().NumField(); i++ { 174 child := sv.Type().Field(i) 175 v := sv.Field(i) 176 // make sure we have a pointer 177 if v.Kind() != reflect.Ptr { 178 v = v.Addr() 179 } 180 // check if that field is a flag or contains flags 181 if fp := addFlags(f, child, v); fp != nil { 182 p = fp 183 } 184 } 185 return p 186 } 187 switch v := value.Interface().(type) { 188 case flag.Value: 189 f.Var(v, flagName, help) 190 case *bool: 191 f.BoolVar(v, flagName, *v, help) 192 case *time.Duration: 193 f.DurationVar(v, flagName, *v, help) 194 case *float64: 195 f.Float64Var(v, flagName, *v, help) 196 case *int64: 197 f.Int64Var(v, flagName, *v, help) 198 case *int: 199 f.IntVar(v, flagName, *v, help) 200 case *string: 201 f.StringVar(v, flagName, *v, help) 202 case *uint: 203 f.UintVar(v, flagName, *v, help) 204 case *uint64: 205 f.Uint64Var(v, flagName, *v, help) 206 default: 207 log.Fatalf("Cannot understand flag of type %T", v) 208 } 209 return nil 210 }