internal/tool/tool.go - tools.git - Git at Google

 // Copyright 2018 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.

 // Package tool is a harness for writing Go tools.
 package tool

 import (
 	"context"
 	"flag"
 	"fmt"
 	"log"
 	"os"
 	"reflect"
 	"runtime"
 	"runtime/pprof"
 	"runtime/trace"
 	"strings"
 	"time"
 )

 // This file is a harness for writing your main function.
 // The original version of the file is in golang.org/x/tools/internal/tool.
 //
 // It adds a method to the Application type
 //     Main(name, usage string, args []string)
 // which should normally be invoked from a true main as follows:
 //     func main() {
 //       (&Application{}).Main("myapp", "non-flag-command-line-arg-help", os.Args[1:])
 //     }
 // It recursively scans the application object for fields with a tag containing
 //     `flag:"flagnames" help:"short help text"`
 // uses all those fields to build command line flags. It will split flagnames on
 // commas and add a flag per name.
 // It expects the Application type to have a method
 //     Run(context.Context, args...string) error
 // which it invokes only after all command line flag processing has been finished.
 // If Run returns an error, the error will be printed to stderr and the
 // application will quit with a non zero exit status.

 // Profile can be embedded in your application struct to automatically
 // add command line arguments and handling for the common profiling methods.
 type Profile struct {
 	CPU    string `flag:"profile.cpu" help:"write CPU profile to this file"`
 	Memory string `flag:"profile.mem" help:"write memory profile to this file"`
 	Alloc  string `flag:"profile.alloc" help:"write alloc profile to this file"`
 	Trace  string `flag:"profile.trace" help:"write trace log to this file"`
 	Block  string `flag:"profile.block" help:"write block profile to this file"`
 }

 // Application is the interface that must be satisfied by an object passed to Main.
 type Application interface {
 	// Name returns the application's name. It is used in help and error messages.
 	Name() string
 	// Most of the help usage is automatically generated, this string should only
 	// describe the contents of non flag arguments.
 	Usage() string
 	// ShortHelp returns the one line overview of the command.
 	ShortHelp() string
 	// DetailedHelp should print a detailed help message. It will only ever be shown
 	// when the ShortHelp is also printed, so there is no need to duplicate
 	// anything from there.
 	// It is passed the flag set so it can print the default values of the flags.
 	// It should use the flag sets configured Output to write the help to.
 	DetailedHelp(*flag.FlagSet)
 	// Run is invoked after all flag processing, and inside the profiling and
 	// error handling harness.
 	Run(ctx context.Context, args ...string) error
 }

 type SubCommand interface {
 	Parent() string
 }

 // This is the type returned by CommandLineErrorf, which causes the outer main
 // to trigger printing of the command line help.
 type commandLineError string

 func (e commandLineError) Error() string { return string(e) }

 // CommandLineErrorf is like fmt.Errorf except that it returns a value that
 // triggers printing of the command line help.
 // In general you should use this when generating command line validation errors.
 func CommandLineErrorf(message string, args ...any) error {
 	return commandLineError(fmt.Sprintf(message, args...))
 }

 // Main should be invoked directly by main function.
 // It will only return if there was no error.  If an error
 // was encountered it is printed to standard error and the
 // application exits with an exit code of 2.
 func Main(ctx context.Context, app Application, args []string) {
 	s := flag.NewFlagSet(app.Name(), flag.ExitOnError)
 	if err := Run(ctx, s, app, args); err != nil {
 		fmt.Fprintf(s.Output(), "%s: %v\n", app.Name(), err)
 		if _, printHelp := err.(commandLineError); printHelp {
 			// TODO(adonovan): refine this. It causes
 			// any command-line error to result in the full
 			// usage message, which typically obscures
 			// the actual error.
 			s.Usage()
 		}
 		os.Exit(2)
 	}
 }

 // Run is the inner loop for Main; invoked by Main, recursively by
 // Run, and by various tests.  It runs the application and returns an
 // error.
 func Run(ctx context.Context, s *flag.FlagSet, app Application, args []string) (resultErr error) {
 	s.Usage = func() {
 		if app.ShortHelp() != "" {
 			fmt.Fprintf(s.Output(), "%s\n\nUsage:\n  ", app.ShortHelp())
 			if sub, ok := app.(SubCommand); ok && sub.Parent() != "" {
 				fmt.Fprintf(s.Output(), "%s [flags] %s", sub.Parent(), app.Name())
 			} else {
 				fmt.Fprintf(s.Output(), "%s [flags]", app.Name())
 			}
 			if usage := app.Usage(); usage != "" {
 				fmt.Fprintf(s.Output(), " %s", usage)
 			}
 			fmt.Fprint(s.Output(), "\n")
 		}
 		app.DetailedHelp(s)
 	}
 	p := addFlags(s, reflect.StructField{}, reflect.ValueOf(app))
 	if err := s.Parse(args); err != nil {
 		return err
 	}

 	if p != nil && p.CPU != "" {
 		f, err := os.Create(p.CPU)
 		if err != nil {
 			return err
 		}
 		if err := pprof.StartCPUProfile(f); err != nil {
 			f.Close() // ignore error
 			return err
 		}
 		defer func() {
 			pprof.StopCPUProfile()
 			if closeErr := f.Close(); resultErr == nil {
 				resultErr = closeErr
 			}
 		}()
 	}

 	if p != nil && p.Trace != "" {
 		f, err := os.Create(p.Trace)
 		if err != nil {
 			return err
 		}
 		if err := trace.Start(f); err != nil {
 			f.Close() // ignore error
 			return err
 		}
 		defer func() {
 			trace.Stop()
 			if closeErr := f.Close(); resultErr == nil {
 				resultErr = closeErr
 			}
 			log.Printf("To view the trace, run:\n$ go tool trace view %s", p.Trace)
 		}()
 	}

 	if p != nil && p.Memory != "" {
 		f, err := os.Create(p.Memory)
 		if err != nil {
 			return err
 		}
 		defer func() {
 			runtime.GC() // get up-to-date statistics
 			if err := pprof.WriteHeapProfile(f); err != nil {
 				log.Printf("Writing memory profile: %v", err)
 			}
 			if err := f.Close(); err != nil {
 				log.Printf("Closing memory profile: %v", err)
 			}
 		}()
 	}

 	if p != nil && p.Alloc != "" {
 		f, err := os.Create(p.Alloc)
 		if err != nil {
 			return err
 		}
 		defer func() {
 			if err := pprof.Lookup("allocs").WriteTo(f, 0); err != nil {
 				log.Printf("Writing alloc profile: %v", err)
 			}
 			if err := f.Close(); err != nil {
 				log.Printf("Closing alloc profile: %v", err)
 			}
 		}()
 	}

 	if p != nil && p.Block != "" {
 		f, err := os.Create(p.Block)
 		if err != nil {
 			return err
 		}
 		runtime.SetBlockProfileRate(1) // record all blocking events
 		defer func() {
 			if err := pprof.Lookup("block").WriteTo(f, 0); err != nil {
 				log.Printf("Writing block profile: %v", err)
 			}
 			if err := f.Close(); err != nil {
 				log.Printf("Closing block profile: %v", err)
 			}
 		}()
 	}

 	return app.Run(ctx, s.Args()...)
 }

 // addFlags scans fields of structs recursively to find things with flag tags
 // and add them to the flag set.
 func addFlags(f *flag.FlagSet, field reflect.StructField, value reflect.Value) *Profile {
 	// is it a field we are allowed to reflect on?
 	if field.PkgPath != "" {
 		return nil
 	}
 	// now see if is actually a flag
 	flagNames, isFlag := field.Tag.Lookup("flag")
 	help := field.Tag.Get("help")
 	if isFlag {
 		nameList := strings.Split(flagNames, ",")
 		// add the main flag
 		addFlag(f, value, nameList[0], help)
 		if len(nameList) > 1 {
 			// and now add any aliases using the same flag value
 			fv := f.Lookup(nameList[0]).Value
 			for _, flagName := range nameList[1:] {
 				f.Var(fv, flagName, help)
 			}
 		}
 		return nil
 	}
 	// not a flag, but it might be a struct with flags in it
 	value = resolve(value.Elem())
 	if value.Kind() != reflect.Struct {
 		return nil
 	}

 	// TODO(adonovan): there's no need for this special treatment of Profile:
 	// The caller can use f.Lookup("profile.cpu") etc instead.
 	p, _ := value.Addr().Interface().(*Profile)
 	// go through all the fields of the struct
 	for i := 0; i < value.Type().NumField(); i++ {
 		child := value.Type().Field(i)
 		v := value.Field(i)
 		// make sure we have a pointer
 		if v.Kind() != reflect.Pointer {
 			v = v.Addr()
 		}
 		// check if that field is a flag or contains flags
 		if fp := addFlags(f, child, v); fp != nil {
 			p = fp
 		}
 	}
 	return p
 }

 func addFlag(f *flag.FlagSet, value reflect.Value, flagName string, help string) {
 	switch v := value.Interface().(type) {
 	case flag.Value:
 		f.Var(v, flagName, help)
 	case *bool:
 		f.BoolVar(v, flagName, *v, help)
 	case *time.Duration:
 		f.DurationVar(v, flagName, *v, help)
 	case *float64:
 		f.Float64Var(v, flagName, *v, help)
 	case *int64:
 		f.Int64Var(v, flagName, *v, help)
 	case *int:
 		f.IntVar(v, flagName, *v, help)
 	case *string:
 		f.StringVar(v, flagName, *v, help)
 	case *uint:
 		f.UintVar(v, flagName, *v, help)
 	case *uint64:
 		f.Uint64Var(v, flagName, *v, help)
 	default:
 		log.Fatalf("field %q of type %T is not assignable to flag.Value", flagName, v)
 	}
 }

 func resolve(v reflect.Value) reflect.Value {
 	for {
 		switch v.Kind() {
 		case reflect.Interface, reflect.Pointer:
 			v = v.Elem()
 		default:
 			return v
 		}
 	}
 }
	// Copyright 2018 The Go Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style
	// license that can be found in the LICENSE file.

	// Package tool is a harness for writing Go tools.
	package tool

	import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"reflect"
	"runtime"
	"runtime/pprof"
	"runtime/trace"
	"strings"
	"time"
	)

	// This file is a harness for writing your main function.
	// The original version of the file is in golang.org/x/tools/internal/tool.
	//
	// It adds a method to the Application type
	// Main(name, usage string, args []string)
	// which should normally be invoked from a true main as follows:
	// func main() {
	// (&Application{}).Main("myapp", "non-flag-command-line-arg-help", os.Args[1:])
	// }
	// It recursively scans the application object for fields with a tag containing
	// `flag:"flagnames" help:"short help text"`
	// uses all those fields to build command line flags. It will split flagnames on
	// commas and add a flag per name.
	// It expects the Application type to have a method
	// Run(context.Context, args...string) error
	// which it invokes only after all command line flag processing has been finished.
	// If Run returns an error, the error will be printed to stderr and the
	// application will quit with a non zero exit status.

	// Profile can be embedded in your application struct to automatically
	// add command line arguments and handling for the common profiling methods.
	type Profile struct {
	CPU string `flag:"profile.cpu" help:"write CPU profile to this file"`
	Memory string `flag:"profile.mem" help:"write memory profile to this file"`
	Alloc string `flag:"profile.alloc" help:"write alloc profile to this file"`
	Trace string `flag:"profile.trace" help:"write trace log to this file"`
	Block string `flag:"profile.block" help:"write block profile to this file"`
	}

	// Application is the interface that must be satisfied by an object passed to Main.
	type Application interface {
	// Name returns the application's name. It is used in help and error messages.
	Name() string
	// Most of the help usage is automatically generated, this string should only
	// describe the contents of non flag arguments.
	Usage() string
	// ShortHelp returns the one line overview of the command.
	ShortHelp() string
	// DetailedHelp should print a detailed help message. It will only ever be shown
	// when the ShortHelp is also printed, so there is no need to duplicate
	// anything from there.
	// It is passed the flag set so it can print the default values of the flags.
	// It should use the flag sets configured Output to write the help to.
	DetailedHelp(*flag.FlagSet)
	// Run is invoked after all flag processing, and inside the profiling and
	// error handling harness.
	Run(ctx context.Context, args ...string) error
	}

	type SubCommand interface {
	Parent() string
	}

	// This is the type returned by CommandLineErrorf, which causes the outer main
	// to trigger printing of the command line help.
	type commandLineError string

	func (e commandLineError) Error() string { return string(e) }

	// CommandLineErrorf is like fmt.Errorf except that it returns a value that
	// triggers printing of the command line help.
	// In general you should use this when generating command line validation errors.
	func CommandLineErrorf(message string, args ...any) error {
	return commandLineError(fmt.Sprintf(message, args...))
	}

	// Main should be invoked directly by main function.
	// It will only return if there was no error. If an error
	// was encountered it is printed to standard error and the
	// application exits with an exit code of 2.
	func Main(ctx context.Context, app Application, args []string) {
	s := flag.NewFlagSet(app.Name(), flag.ExitOnError)
	if err := Run(ctx, s, app, args); err != nil {
	fmt.Fprintf(s.Output(), "%s: %v\n", app.Name(), err)
	if _, printHelp := err.(commandLineError); printHelp {
	// TODO(adonovan): refine this. It causes
	// any command-line error to result in the full
	// usage message, which typically obscures
	// the actual error.
	s.Usage()
	}
	os.Exit(2)
	}
	}

	// Run is the inner loop for Main; invoked by Main, recursively by
	// Run, and by various tests. It runs the application and returns an
	// error.
	func Run(ctx context.Context, s *flag.FlagSet, app Application, args []string) (resultErr error) {
	s.Usage = func() {
	if app.ShortHelp() != "" {
	fmt.Fprintf(s.Output(), "%s\n\nUsage:\n ", app.ShortHelp())
	if sub, ok := app.(SubCommand); ok && sub.Parent() != "" {
	fmt.Fprintf(s.Output(), "%s [flags] %s", sub.Parent(), app.Name())
	} else {
	fmt.Fprintf(s.Output(), "%s [flags]", app.Name())
	}
	if usage := app.Usage(); usage != "" {
	fmt.Fprintf(s.Output(), " %s", usage)
	}
	fmt.Fprint(s.Output(), "\n")
	}
	app.DetailedHelp(s)
	}
	p := addFlags(s, reflect.StructField{}, reflect.ValueOf(app))
	if err := s.Parse(args); err != nil {
	return err
	}

	if p != nil && p.CPU != "" {
	f, err := os.Create(p.CPU)
	if err != nil {
	return err
	}
	if err := pprof.StartCPUProfile(f); err != nil {
	f.Close() // ignore error
	return err
	}
	defer func() {
	pprof.StopCPUProfile()
	if closeErr := f.Close(); resultErr == nil {
	resultErr = closeErr
	}
	}()
	}

	if p != nil && p.Trace != "" {
	f, err := os.Create(p.Trace)
	if err != nil {
	return err
	}
	if err := trace.Start(f); err != nil {
	f.Close() // ignore error
	return err
	}
	defer func() {
	trace.Stop()
	if closeErr := f.Close(); resultErr == nil {
	resultErr = closeErr
	}
	log.Printf("To view the trace, run:\n$ go tool trace view %s", p.Trace)
	}()
	}

	if p != nil && p.Memory != "" {
	f, err := os.Create(p.Memory)
	if err != nil {
	return err
	}
	defer func() {
	runtime.GC() // get up-to-date statistics
	if err := pprof.WriteHeapProfile(f); err != nil {
	log.Printf("Writing memory profile: %v", err)
	}
	if err := f.Close(); err != nil {
	log.Printf("Closing memory profile: %v", err)
	}
	}()
	}

	if p != nil && p.Alloc != "" {
	f, err := os.Create(p.Alloc)
	if err != nil {
	return err
	}
	defer func() {
	if err := pprof.Lookup("allocs").WriteTo(f, 0); err != nil {
	log.Printf("Writing alloc profile: %v", err)
	}
	if err := f.Close(); err != nil {
	log.Printf("Closing alloc profile: %v", err)
	}
	}()
	}

	if p != nil && p.Block != "" {
	f, err := os.Create(p.Block)
	if err != nil {
	return err
	}
	runtime.SetBlockProfileRate(1) // record all blocking events
	defer func() {
	if err := pprof.Lookup("block").WriteTo(f, 0); err != nil {
	log.Printf("Writing block profile: %v", err)
	}
	if err := f.Close(); err != nil {
	log.Printf("Closing block profile: %v", err)
	}
	}()
	}

	return app.Run(ctx, s.Args()...)
	}

	// addFlags scans fields of structs recursively to find things with flag tags
	// and add them to the flag set.
	func addFlags(f flag.FlagSet, field reflect.StructField, value reflect.Value) Profile {
	// is it a field we are allowed to reflect on?
	if field.PkgPath != "" {
	return nil
	}
	// now see if is actually a flag
	flagNames, isFlag := field.Tag.Lookup("flag")
	help := field.Tag.Get("help")
	if isFlag {
	nameList := strings.Split(flagNames, ",")
	// add the main flag
	addFlag(f, value, nameList[0], help)
	if len(nameList) > 1 {
	// and now add any aliases using the same flag value
	fv := f.Lookup(nameList[0]).Value
	for _, flagName := range nameList[1:] {
	f.Var(fv, flagName, help)
	}
	}
	return nil
	}
	// not a flag, but it might be a struct with flags in it
	value = resolve(value.Elem())
	if value.Kind() != reflect.Struct {
	return nil
	}

	// TODO(adonovan): there's no need for this special treatment of Profile:
	// The caller can use f.Lookup("profile.cpu") etc instead.
	p, _ := value.Addr().Interface().(*Profile)
	// go through all the fields of the struct
	for i := 0; i < value.Type().NumField(); i++ {
	child := value.Type().Field(i)
	v := value.Field(i)
	// make sure we have a pointer
	if v.Kind() != reflect.Pointer {
	v = v.Addr()
	}
	// check if that field is a flag or contains flags
	if fp := addFlags(f, child, v); fp != nil {
	p = fp
	}
	}
	return p
	}

	func addFlag(f *flag.FlagSet, value reflect.Value, flagName string, help string) {
	switch v := value.Interface().(type) {
	case flag.Value:
	f.Var(v, flagName, help)
	case *bool:
	f.BoolVar(v, flagName, *v, help)
	case *time.Duration:
	f.DurationVar(v, flagName, *v, help)
	case *float64:
	f.Float64Var(v, flagName, *v, help)
	case *int64:
	f.Int64Var(v, flagName, *v, help)
	case *int:
	f.IntVar(v, flagName, *v, help)
	case *string:
	f.StringVar(v, flagName, *v, help)
	case *uint:
	f.UintVar(v, flagName, *v, help)
	case *uint64:
	f.Uint64Var(v, flagName, *v, help)
	default:
	log.Fatalf("field %q of type %T is not assignable to flag.Value", flagName, v)
	}
	}

	func resolve(v reflect.Value) reflect.Value {
	for {
	switch v.Kind() {
	case reflect.Interface, reflect.Pointer:
	v = v.Elem()
	default:
	return v
	}
	}
	}