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