tool.go

Documentation: golang.org/x/tools/internal/tool

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