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 `