lab.nexedi.com/kirr/go123@v0.0.0-20240207185015-8299741fa871/prog/prog.go (about)

     1  // Copyright (C) 2017-2019  Nexedi SA and Contributors.
     2  //                          Kirill Smelkov <kirr@nexedi.com>
     3  //
     4  // This program is free software: you can Use, Study, Modify and Redistribute
     5  // it under the terms of the GNU General Public License version 3, or (at your
     6  // option) any later version, as published by the Free Software Foundation.
     7  //
     8  // You can also Link and Combine this program with other software covered by
     9  // the terms of any of the Free Software licenses or any of the Open Source
    10  // Initiative approved licenses and Convey the resulting work. Corresponding
    11  // source of such a combination shall include the source code for all other
    12  // software used.
    13  //
    14  // This program is distributed WITHOUT ANY WARRANTY; without even the implied
    15  // warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
    16  //
    17  // See COPYING file for full licensing terms.
    18  // See https://www.nexedi.com/licensing for rationale and options.
    19  
    20  // Package prog provides infrastructure for implementing programs with subcommands.
    21  //
    22  // Usage is simple: initialize a MainProg var appropriately and call its .Main().
    23  package prog
    24  
    25  import (
    26  	"flag"
    27  	"fmt"
    28  	"io"
    29  	"log"
    30  	"os"
    31  	"runtime"
    32  	"runtime/pprof"
    33  	"runtime/trace"
    34  )
    35  
    36  // Command describes one program subcommand.
    37  type Command struct {
    38  	Name    string
    39  	Summary string
    40  	Usage   func(w io.Writer)
    41  	Main    func(argv []string)
    42  }
    43  
    44  // CommandRegistry is ordered collection of Commands.
    45  type CommandRegistry []Command
    46  
    47  // Lookup returns Command with corresponding name or nil.
    48  func (cmdv CommandRegistry) Lookup(command string) *Command {
    49  	for i := range cmdv {
    50  		if cmdv[i].Name == command {
    51  			return &cmdv[i]
    52  		}
    53  	}
    54  	return nil
    55  }
    56  
    57  // HelpTopic describes one help topic.
    58  type HelpTopic struct {
    59  	Name    string
    60  	Summary string
    61  	Text    string
    62  }
    63  
    64  // HelpRegistry is ordered collection of HelpTopics.
    65  type HelpRegistry []HelpTopic
    66  
    67  // Lookup returns HelpTopic with corresponding name or nil.
    68  func (helpv HelpRegistry) Lookup(topic string) *HelpTopic {
    69  	for i := range helpv {
    70  		if helpv[i].Name == topic {
    71  			return &helpv[i]
    72  		}
    73  	}
    74  	return nil
    75  }
    76  
    77  // ----------------------------------------
    78  
    79  // MainProg defines a program to run with subcommands and help topics.
    80  type MainProg struct {
    81  	Name       string          // name of the program, e.g. "zodb"
    82  	Summary    string          // 1-line summary of what program does
    83  	Commands   CommandRegistry // provided subcommands
    84  	HelpTopics HelpRegistry    // provided help topics
    85  }
    86  
    87  // Exit is like os.Exit but makes sure deferred functions are run.
    88  //
    89  // Exit should be called from main goroutine.
    90  func Exit(code int) {
    91  	panic(&programExit{code})
    92  }
    93  
    94  // Fatal is like log.Fatal but makes sure deferred functions are run.
    95  //
    96  // Fatal should be called from main goroutine.
    97  func Fatal(v ...interface{}) {
    98  	log.Print(v...)
    99  	Exit(1)
   100  }
   101  
   102  // programExit is thrown when Exit or Fatal are called.
   103  type programExit struct {
   104  	code int
   105  }
   106  
   107  // Main is the main entry point for the program. Call it from main.
   108  //
   109  // Do not call os.Exit or log.Fatal from your program. Instead use Exit and
   110  // Fatal from prog package so that deferred functions setup by Main could
   111  // be run.
   112  func (prog *MainProg) Main() {
   113  	// handle exit throw-requests
   114  	defer func() {
   115  		r := recover()
   116  		if e, _ := r.(*programExit); e != nil {
   117  			// TODO log.Flush()
   118  			os.Exit(e.code)
   119  		}
   120  		if r != nil {
   121  			panic(r)
   122  		}
   123  	}()
   124  
   125  	prog.main()
   126  }
   127  
   128  func (prog *MainProg) main() {
   129  	flag.Usage = prog.usage
   130  	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to `file`")
   131  	memprofile := flag.String("memprofile", "", "write memory profile to `file`")
   132  	traceout := flag.String("trace", "", "write execution trace to `file`")
   133  	flag.Parse()
   134  	argv := flag.Args()
   135  
   136  	if len(argv) == 0 {
   137  		prog.usage()
   138  		Exit(2)
   139  	}
   140  
   141  	command := argv[0]
   142  
   143  	// handle common options
   144  	if *cpuprofile != "" {
   145  		f, err := os.Create(*cpuprofile)
   146  		if err != nil {
   147  			Fatal("could not create CPU profile: ", err)
   148  		}
   149  		if err := pprof.StartCPUProfile(f); err != nil {
   150  			Fatal("could not start CPU profile: ", err)
   151  		}
   152  		defer pprof.StopCPUProfile()
   153  	}
   154  
   155  	defer func() {
   156  		if *memprofile != "" {
   157  			f, err := os.Create(*memprofile)
   158  			if err != nil {
   159  				Fatal("could not create memory profile: ", err)
   160  			}
   161  			runtime.GC() // get up-to-date statistics
   162  			if err := pprof.WriteHeapProfile(f); err != nil {
   163  				Fatal("could not write memory profile: ", err)
   164  			}
   165  			f.Close()
   166  		}
   167  	}()
   168  
   169  	if *traceout != "" {
   170  		f, err := os.Create(*traceout)
   171  		if err != nil {
   172  			Fatal("could not create trace: ", err)
   173  		}
   174  		defer func() {
   175  			if err := f.Close(); err != nil {
   176  				Fatal("could not close trace: ", err)
   177  			}
   178  		}()
   179  		if err := trace.Start(f); err != nil {
   180  			Fatal("could not start trace: ", err)
   181  		}
   182  		defer trace.Stop()
   183  	}
   184  
   185  
   186  	// help on a topic
   187  	if command == "help" {
   188  		prog.help(argv)
   189  		return
   190  	}
   191  
   192  	// run subcommand
   193  	cmd := prog.Commands.Lookup(command)
   194  	if cmd == nil {
   195  		fmt.Fprintf(os.Stderr, "%s: unknown subcommand \"%s\"\n", prog.Name, command)
   196  		fmt.Fprintf(os.Stderr, "Run '%s help' for usage.\n", prog.Name)
   197  		Exit(2)
   198  	}
   199  
   200  	cmd.Main(argv)
   201  }
   202  
   203  // usage shows usage text for whole program.
   204  func (prog *MainProg) usage() {
   205  	w := os.Stderr
   206  	fmt.Fprintf(w,
   207  `%s.
   208  
   209  Usage:
   210  
   211  	%s [options] command [arguments]
   212  
   213  The commands are:
   214  
   215  `, prog.Summary, prog.Name)
   216  
   217  	// to lalign commands & help summaries
   218  	nameWidth := 0
   219  	for _, cmd := range prog.Commands {
   220  		if len(cmd.Name) > nameWidth {
   221  			nameWidth = len(cmd.Name)
   222  		}
   223  	}
   224  	for _, topic := range prog.helpTopics() {
   225  		if len(topic.Name) > nameWidth {
   226  			nameWidth = len(topic.Name)
   227  		}
   228  	}
   229  
   230  	for _, cmd := range prog.Commands {
   231  		fmt.Fprintf(w, "\t%-*s %s\n", nameWidth, cmd.Name, cmd.Summary)
   232  	}
   233  
   234  	fmt.Fprintf(w,
   235  `
   236  
   237  Use "%s help [command]" for more information about a command.
   238  `, prog.Name)
   239  
   240  	if len(prog.helpTopics()) > 0 {
   241  		fmt.Fprintf(w,
   242  `
   243  Additional help topics:
   244  
   245  `)
   246  
   247  		for _, topic := range prog.helpTopics() {
   248  			fmt.Fprintf(w, "\t%-*s %s\n", nameWidth, topic.Name, topic.Summary)
   249  		}
   250  
   251  		fmt.Fprintf(w,
   252  `
   253  Use "%s help [topic]" for more information about that topic.
   254  
   255  `, prog.Name)
   256  	}
   257  }
   258  
   259  
   260  // help shows general help or help for a command/topic.
   261  func (prog *MainProg) help(argv []string) {
   262  	if len(argv) < 2 {	// help topic ...
   263  		prog.usage()
   264  		Exit(2)
   265  	}
   266  
   267  	topic := argv[1]
   268  
   269  	// topic can either be a command name or a help topic
   270  	command := prog.Commands.Lookup(topic)
   271  	if command != nil {
   272  		command.Usage(os.Stdout)
   273  		Exit(0)
   274  	}
   275  
   276  	helpTopic := prog.helpTopics().Lookup(topic)
   277  	if helpTopic != nil {
   278  		fmt.Println(helpTopic.Text)
   279  		Exit(0)
   280  	}
   281  
   282  	fmt.Fprintf(os.Stderr, "Unknown help topic `%s`.  Run '%s help'.\n", topic, prog.Name)
   283  	Exit(2)
   284  }
   285  
   286  // helpTopics returns provided help topics augmented with help on common topics
   287  // provided by prog driver.
   288  func (prog *MainProg) helpTopics() HelpRegistry {
   289  	return append(helpCommon, prog.HelpTopics...)
   290  }
   291  
   292  var helpCommon = HelpRegistry{
   293  	{"options", "options common to all commands", helpOptions},
   294  }
   295  
   296  const helpOptions =
   297  `Options common to all commands:
   298  
   299  	-cpuprofile <file>	write cpu profile to <file>
   300  	-memprofile <file>	write memory profile to <file>
   301  
   302  	TODO also document glog options
   303  `