Merge pull request #86 from xushiwei/q

x/cmdline/app

cmdline/app/app.go

@@ -0,0 +1,288 @@
+/*
+ Copyright 2023 Qiniu Limited (qiniu.com)
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+*/
+
+package app
+
+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"`
+}
+
+// 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 ...interface{}) 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)
+			}
+			f.Close()
+		}()
+	}
+
+	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)
+			}
+			f.Close()
+		}()
+	}
+
+	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.Ptr {
+			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("Cannot understand flag of type %T", v)
+	}
+}
+
+func resolve(v reflect.Value) reflect.Value {
+	for {
+		switch v.Kind() {
+		case reflect.Interface, reflect.Ptr:
+			v = v.Elem()
+		default:
+			return v
+		}
+	}
+}