From 9069aa4d19af2268603e659926f9b868e817516c Mon Sep 17 00:00:00 2001 From: Glib Smaga Date: Tue, 28 Jun 2022 09:18:41 -0700 Subject: [PATCH] Migrate gops to use spf13/cobra for command execution Manual command setup was starting to show some strain in terms of documentation and also adding new features (regardless of what the protocol supports currently). This new setup aims to make it easier to add new documentation and functionality. In comparison to the previous version, this increased the binary size by 2.4M. ``` 6.4M gops 4.0M gops_master ``` --- go.mod | 1 + go.sum | 9 +++ internal/cmd/process.go | 32 +++++++++++ internal/cmd/root.go | 100 ++++---------------------------- internal/cmd/shared.go | 111 ++++++++++++++++++++++++++++++++---- internal/cmd/shared_test.go | 39 +++++++++++++ internal/cmd/tree.go | 12 ++++ main.go | 27 ++++++++- 8 files changed, 232 insertions(+), 99 deletions(-) create mode 100644 internal/cmd/shared_test.go diff --git a/go.mod b/go.mod index cf8fd4b..eec5642 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,7 @@ go 1.13 require ( github.com/keybase/go-ps v0.0.0-20190827175125-91aafc93ba19 github.com/shirou/gopsutil/v3 v3.22.4 + github.com/spf13/cobra v1.4.0 github.com/xlab/treeprint v1.1.0 golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a rsc.io/goversion v1.2.0 diff --git a/go.sum b/go.sum index 98b8f6b..a135bd8 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,4 @@ +github.com/cpuguy83/go-md2man/v2 v2.0.1/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o= github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY= @@ -5,6 +6,8 @@ github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiU github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.7 h1:81/ik6ipDQS2aGcBfIN5dHDB36BwrStyeAQquSYCV4o= github.com/google/go-cmp v0.5.7/go.mod h1:n+brtR0CgQNWTVd5ZUFpTBC8YFBDLK/h/bpaJ8/DtOE= +github.com/inconshreveable/mousetrap v1.0.0 h1:Z8tu5sraLXCXIcARxBp/8cbvlwVa7Z1NHg9XEKhtSvM= +github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8= github.com/keybase/go-ps v0.0.0-20190827175125-91aafc93ba19 h1:WjT3fLi9n8YWh/Ih8Q1LHAPsTqGddPcHqscN+PJ3i68= github.com/keybase/go-ps v0.0.0-20190827175125-91aafc93ba19/go.mod h1:hY+WOq6m2FpbvyrI93sMaypsttvaIL5nhVR92dTMUcQ= github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4= @@ -13,8 +16,13 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw= github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE= +github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= github.com/shirou/gopsutil/v3 v3.22.4 h1:srAQaiX6jX/cYL6q29aE0m8lOskT9CurZ9N61YR3yoI= github.com/shirou/gopsutil/v3 v3.22.4/go.mod h1:D01hZJ4pVHPpCTZ3m3T2+wDF2YAGfd+H4ifUguaQzHM= +github.com/spf13/cobra v1.4.0 h1:y+wJpx64xcgO1V+RcnwW0LEHxTKRi2ZDPSBjWnrg88Q= +github.com/spf13/cobra v1.4.0/go.mod h1:Wo4iy3BUC+X2Fybo0PDqwJIv3dNRiZLHQymsfxlB84g= +github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA= +github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY= @@ -36,6 +44,7 @@ golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IV golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= rsc.io/goversion v1.2.0 h1:SPn+NLTiAG7w30IRK/DKp1BjvpWabYgxlLp/+kx5J8w= diff --git a/internal/cmd/process.go b/internal/cmd/process.go index d4488ca..bb5d49b 100644 --- a/internal/cmd/process.go +++ b/internal/cmd/process.go @@ -8,12 +8,44 @@ import ( "fmt" "log" "math" + "strconv" "strings" "time" "github.com/shirou/gopsutil/v3/process" + "github.com/spf13/cobra" ) +// ProcessCommand displays information about a Go process. +func ProcessCommand() *cobra.Command { + return &cobra.Command{ + Use: "process", + Aliases: []string{"pid", "proc"}, + Short: "Prints information about a Go process.", + RunE: func(cmd *cobra.Command, args []string) error { + ProcessInfo(args) + return nil + }, + } +} + +// ProcessInfo takes arguments starting with pid|:addr and grabs all kinds of +// useful Go process information. +func ProcessInfo(args []string) { + pid, err := strconv.Atoi(args[0]) + + var period time.Duration + if len(args) >= 2 { + period, err = time.ParseDuration(args[1]) + if err != nil { + secs, _ := strconv.Atoi(args[1]) + period = time.Duration(secs) * time.Second + } + } + + processInfo(pid, period) +} + func processInfo(pid int, period time.Duration) { if period < 0 { log.Fatalf("Cannot determine CPU usage for negative duration %v", period) diff --git a/internal/cmd/root.go b/internal/cmd/root.go index 8496004..fb50bf0 100644 --- a/internal/cmd/root.go +++ b/internal/cmd/root.go @@ -11,87 +11,23 @@ import ( "regexp" "strconv" "strings" - "time" "github.com/google/gops/goprocess" + "github.com/spf13/cobra" ) -const helpText = `gops is a tool to list and diagnose Go processes. - -Usage: - gops ... +// NewRoot command. +func NewRoot() *cobra.Command { + return &cobra.Command{ + Use: "gops", + Short: "gops is a tool to list and diagnose Go processes.", + Example: ` gops ... gops # displays process info - gops help # displays this help message - -Commands: - stack Prints the stack trace. - gc Runs the garbage collector and blocks until successful. - setgc Sets the garbage collection target percentage. - memstats Prints the allocation and garbage collection stats. - version Prints the Go version used to build the program. - stats Prints runtime stats. - trace Runs the runtime tracer for 5 secs and launches "go tool trace". - pprof-heap Reads the heap profile and launches "go tool pprof". - pprof-cpu Reads the CPU profile and launches "go tool pprof". - -All commands require the agent running on the Go process. -"*" indicates the process is running the agent.` - -// TODO(jbd): add link that explains the use of agent. - -// Execute the root command. -func Execute() { - if len(os.Args) < 2 { - processes() - return - } - - cmd := os.Args[1] - - // See if it is a PID. - pid, err := strconv.Atoi(cmd) - if err == nil { - var period time.Duration - if len(os.Args) >= 3 { - period, err = time.ParseDuration(os.Args[2]) - if err != nil { - secs, _ := strconv.Atoi(os.Args[2]) - period = time.Duration(secs) * time.Second - } - } - processInfo(pid, period) - return - } - - if cmd == "help" { - usage("") - } - - if cmd == "tree" { - displayProcessTree() - return - } - - fn, ok := cmds[cmd] - if !ok { - usage("unknown subcommand") - } - if len(os.Args) < 3 { - usage("Missing PID or address.") - os.Exit(1) - } - addr, err := targetToAddr(os.Args[2]) - if err != nil { - fmt.Fprintf(os.Stderr, "Couldn't resolve addr or pid %v to TCPAddress: %v\n", os.Args[2], err) - os.Exit(1) - } - var params []string - if len(os.Args) > 3 { - params = append(params, os.Args[3:]...) - } - if err := fn(*addr, params); err != nil { - fmt.Fprintf(os.Stderr, "%v\n", err) - os.Exit(1) + gops help # displays this help message`, + // TODO(jbd): add link that explains the use of agent. + Run: func(cmd *cobra.Command, args []string) { + processes() + }, } } @@ -144,18 +80,6 @@ func shortenVersion(v string) string { return results[0] } -func usage(msg string) { - // default exit code for the statement - exitCode := 0 - if msg != "" { - // founded an unexpected command - fmt.Printf("gops: %v\n", msg) - exitCode = 1 - } - fmt.Fprintf(os.Stderr, "%v\n", helpText) - os.Exit(exitCode) -} - func pad(s string, total int) string { if len(s) >= total { return s diff --git a/internal/cmd/shared.go b/internal/cmd/shared.go index 9f66fcf..c90880d 100644 --- a/internal/cmd/shared.go +++ b/internal/cmd/shared.go @@ -18,18 +18,109 @@ import ( "github.com/google/gops/internal" "github.com/google/gops/signal" + "github.com/spf13/cobra" ) -var cmds = map[string](func(addr net.TCPAddr, params []string) error){ - "stack": stackTrace, - "gc": gc, - "memstats": memStats, - "version": version, - "pprof-heap": pprofHeap, - "pprof-cpu": pprofCPU, - "stats": stats, - "trace": trace, - "setgc": setGC, +// AgentCommands is a bridge between the legacy multiplexing to commands, and +// full migration to Cobra for each command. +// +// The code is already nicely structured with one function per command so it +// seemed cleaner to combine them all together here and "generate" cobra +// commands as just thin wrappers, rather through individual constructors. +func AgentCommands() []*cobra.Command { + var res []*cobra.Command + + var cmds = []legacyCommand{ + { + name: "stack", + short: "Prints the stack trace.", + fn: stackTrace, + }, + { + name: "gc", + short: "Runs the garbage collector and blocks until successful.", + fn: gc, + }, + { + name: "setgc", + short: "Sets the garbage collection target percentage.", + fn: setGC, + }, + { + name: "memstats", + short: "Prints the allocation and garbage collection stats.", + fn: memStats, + }, + { + name: "stats", + short: "Prints runtime stats.", + fn: stats, + }, + { + name: "trace", + short: "Runs the runtime tracer for 5 secs and launches \"go tool trace\".", + fn: trace, + }, + { + name: "pprof-heap", + short: "Reads the heap profile and launches \"go tool pprof\".", + fn: pprofHeap, + }, + { + name: "pprof-cpu", + short: "Reads the CPU profile and launches \"go tool pprof\".", + fn: pprofCPU, + }, + { + name: "version", + short: "Prints the Go version used to build the program.", + fn: version, + }, + } + + for _, c := range cmds { + c := c + res = append(res, &cobra.Command{ + Use: fmt.Sprintf("%s ", c.name), + Short: c.short, + + RunE: func(cmd *cobra.Command, args []string) error { + if len(args) < 1 { + return fmt.Errorf("missing PID or address") + } + + addr, err := targetToAddr(args[0]) + if err != nil { + return fmt.Errorf( + "couldn't resolve addr or pid %v to TCPAddress: %v\n", args[0], err, + ) + } + + var params []string + if len(args) > 1 { + params = append(params, args[1:]...) + } + + if err := c.fn(*addr, params); err != nil { + return err + } + + return nil + }, + + // errors get double printed otherwise + SilenceUsage: true, + SilenceErrors: true, + }) + } + + return res +} + +type legacyCommand struct { + name string + short string + fn func(addr net.TCPAddr, params []string) error } func setGC(addr net.TCPAddr, params []string) error { diff --git a/internal/cmd/shared_test.go b/internal/cmd/shared_test.go new file mode 100644 index 0000000..a49c6f0 --- /dev/null +++ b/internal/cmd/shared_test.go @@ -0,0 +1,39 @@ +// Copyright 2022 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 cmd + +import ( + "bytes" + "strings" + "testing" + + "github.com/spf13/cobra" +) + +func TestCommandPresence(t *testing.T) { + cmd := &cobra.Command{Use: "gops"} + cmd.AddCommand(AgentCommands()...) + + var out bytes.Buffer + cmd.SetOut(&out) + cmd.SetArgs([]string{"--help"}) + if err := cmd.Execute(); err != nil { + t.Error(err) + } + + // basic check to make sure all the legacy commands have been ported over + // it doesn't test they are correctly _implemented_, just that they are not + // missing. + wants := []string{ + "completion", "gc", "memstats", "pprof-cpu", "pprof-heap", "setgc", + "stack", "stats", "trace", "version", + } + outs := out.String() + for _, want := range wants { + if !strings.Contains(outs, want) { + t.Errorf("%q command not found in help", want) + } + } +} diff --git a/internal/cmd/tree.go b/internal/cmd/tree.go index d814d9a..12c116f 100644 --- a/internal/cmd/tree.go +++ b/internal/cmd/tree.go @@ -10,9 +10,21 @@ import ( "strconv" "github.com/google/gops/goprocess" + "github.com/spf13/cobra" "github.com/xlab/treeprint" ) +// TreeCommand displays a process tree. +func TreeCommand() *cobra.Command { + return &cobra.Command{ + Use: "tree", + Short: "Display parent-child tree for Go processes.", + Run: func(cmd *cobra.Command, args []string) { + displayProcessTree() + }, + } +} + // displayProcessTree displays a tree of all the running Go processes. func displayProcessTree() { ps := goprocess.FindAll() diff --git a/main.go b/main.go index b2cfd89..2a07994 100644 --- a/main.go +++ b/main.go @@ -6,9 +6,34 @@ package main import ( + "log" + "os" + "strconv" + "github.com/google/gops/internal/cmd" ) func main() { - cmd.Execute() + var root = cmd.NewRoot() + root.AddCommand(cmd.ProcessCommand()) + root.AddCommand(cmd.TreeCommand()) + root.AddCommand(cmd.AgentCommands()...) + + // Legacy support for `gops ` command. + // + // When the second argument is provided as int as opposed to a sub-command + // (like proc, version, etc), gops command effectively shortcuts that + // to `gops process `. + if len(os.Args) > 1 { + // See second argument appears to be a pid rather than a subcommand + _, err := strconv.Atoi(os.Args[1]) + if err == nil { + cmd.ProcessInfo(os.Args[1:]) // shift off the command name + return + } + } + + if err := root.Execute(); err != nil { + log.Fatal(err) + } }