diff --git a/main.go b/main.go index a16d1a0..a727bd3 100644 --- a/main.go +++ b/main.go @@ -53,7 +53,9 @@ func main() { func runWithArgs(args ...string) error { opts := &options{} - if _, err := gflags.ParseArgs(opts, args); err != nil { + + remaining, err := gflags.ParseArgs(opts, args) + if err != nil { if flagErr, ok := err.(*gflags.Error); ok && flagErr.Type == gflags.ErrHelp { os.Exit(0) } @@ -63,11 +65,15 @@ func runWithArgs(args ...string) error { return fmt.Errorf("invalid options: %v", err) } - return runWithOptions(opts) + // If there are remaining arguments, the first argument is the binary name. + if len(remaining) > 0 { + remaining = remaining[1:] + } + return runWithOptions(opts, remaining) } -func runWithOptions(opts *options) error { - pprofRawOutput, err := pprof.GetRaw(opts.PProfOptions) +func runWithOptions(opts *options, remaining []string) error { + pprofRawOutput, err := pprof.GetRaw(opts.PProfOptions, remaining) if err != nil { return fmt.Errorf("could not get raw output from pprof: %v", err) } diff --git a/main_test.go b/main_test.go index 32977de..9d8b4c4 100644 --- a/main_test.go +++ b/main_test.go @@ -60,7 +60,13 @@ func TestBadArgs(t *testing.T) { } func TestMain(t *testing.T) { - os.Args = []string{"--raw", "--binaryinput", testPProfInputFile} + os.Args = []string{"go-torch", "--raw", "--binaryinput", testPProfInputFile} + main() + // Test should not fatal. +} + +func TestMainRemaining(t *testing.T) { + os.Args = []string{"go-torch", "--raw", testPProfInputFile} main() // Test should not fatal. } @@ -97,7 +103,7 @@ func TestRunRaw(t *testing.T) { opts := getDefaultOptions() opts.Raw = true - if err := runWithOptions(opts); err != nil { + if err := runWithOptions(opts, nil); err != nil { t.Fatalf("Run with Raw failed: %v", err) } } @@ -117,7 +123,7 @@ func TestRunFile(t *testing.T) { opts.File = getTempFilename(t, ".svg") withScriptsInPath(t, func() { - if err := runWithOptions(opts); err != nil { + if err := runWithOptions(opts, nil); err != nil { t.Fatalf("Run with Print failed: %v", err) } @@ -144,7 +150,7 @@ func TestRunBadFile(t *testing.T) { opts.File = "/dev/zero/invalid/file" withScriptsInPath(t, func() { - if err := runWithOptions(opts); err == nil { + if err := runWithOptions(opts, nil); err == nil { t.Fatalf("Run with bad file expected to fail") } }) @@ -155,7 +161,7 @@ func TestRunPrint(t *testing.T) { opts.Print = true withScriptsInPath(t, func() { - if err := runWithOptions(opts); err != nil { + if err := runWithOptions(opts, nil); err != nil { t.Fatalf("Run with Print failed: %v", err) } // TODO(prashantv): Verify that output is printed to stdout. diff --git a/pprof/pprof.go b/pprof/pprof.go index fe1bc01..311fbd9 100644 --- a/pprof/pprof.go +++ b/pprof/pprof.go @@ -36,13 +36,14 @@ type Options struct { URLSuffix string `short:"s" long:"suffix" default:"/debug/pprof/profile" description:"URL path of pprof profile"` BinaryFile string `short:"b" long:"binaryinput" description:"File path of previously saved binary profile. (binary profile is anything accepted by https://golang.org/cmd/pprof)"` BinaryName string `long:"binaryname" description:"File path of the binary that the binaryinput is for, used for pprof inputs"` - TimeSeconds int `short:"t" long:"time" default:"30" description:"Duration to profile for"` + TimeSeconds int `short:"t" long:"seconds" default:"30" description:"Number of seconds to profile for"` ExtraArgs []string `long:"pprofArgs" description:"Extra arguments for pprof"` + TimeAlias *int `hidden:"true" long:"time" description:"Alias for backwards compatibility"` } // GetRaw returns the raw output from pprof for the given options. -func GetRaw(opts Options) ([]byte, error) { - args, err := getArgs(opts) +func GetRaw(opts Options, remaining []string) ([]byte, error) { + args, err := getArgs(opts, remaining) if err != nil { return nil, err } @@ -51,7 +52,19 @@ func GetRaw(opts Options) ([]byte, error) { } // getArgs gets the arguments to run pprof with for a given set of Options. -func getArgs(opts Options) ([]string, error) { +func getArgs(opts Options, remaining []string) ([]string, error) { + if opts.TimeAlias != nil { + opts.TimeSeconds = *opts.TimeAlias + } + if len(remaining) > 0 { + var pprofArgs []string + if opts.TimeSeconds > 0 { + pprofArgs = append(pprofArgs, "-seconds", fmt.Sprint(opts.TimeSeconds)) + } + pprofArgs = append(pprofArgs, remaining...) + return pprofArgs, nil + } + pprofArgs := opts.ExtraArgs if opts.BinaryFile != "" { if opts.BinaryName != "" { diff --git a/pprof/pprof_test.go b/pprof/pprof_test.go index 35373f3..20d1566 100644 --- a/pprof/pprof_test.go +++ b/pprof/pprof_test.go @@ -22,15 +22,19 @@ package pprof import ( "bytes" + "net/http" + "net/http/httptest" "reflect" "testing" ) func TestGetArgs(t *testing.T) { + four := 4 tests := []struct { - opts Options - expected []string - wantErr bool + opts Options + remaining []string + expected []string + wantErr bool }{ { opts: Options{ @@ -48,6 +52,14 @@ func TestGetArgs(t *testing.T) { }, expected: []string{"-seconds", "5", "http://localhost:1234/path/to/profile"}, }, + { + opts: Options{ + BaseURL: "http://localhost:1234/", + URLSuffix: "/path/to/profile", + TimeAlias: &four, + }, + expected: []string{"-seconds", "4", "http://localhost:1234/path/to/profile"}, + }, { opts: Options{ BaseURL: "http://localhost:1234/test", @@ -90,10 +102,32 @@ func TestGetArgs(t *testing.T) { }, wantErr: true, }, + { + remaining: []string{"binary", "input"}, + expected: []string{"binary", "input"}, + }, + { + opts: Options{ + TimeSeconds: 5, + }, + remaining: []string{"binary", "input"}, + expected: []string{"-seconds", "5", "binary", "input"}, + }, + { + opts: Options{ + TimeSeconds: 5, + // All other fields are ignored when remaining is specified. + BinaryFile: "/path/to/binaryfile", + BinaryName: "/path/to/binaryname", + URLSuffix: "/ignored", + }, + remaining: []string{"binary", "input"}, + expected: []string{"-seconds", "5", "binary", "input"}, + }, } for _, tt := range tests { - got, err := getArgs(tt.opts) + got, err := getArgs(tt.opts, tt.remaining) if (err != nil) != tt.wantErr { t.Errorf("wantErr %v got error: %v", tt.wantErr, err) continue @@ -121,7 +155,10 @@ func TestRunPProfMissingFile(t *testing.T) { } func TestRunPProfInvalidURL(t *testing.T) { - if _, err := runPProf("http://127.0.0.1:999/profile"); err == nil { + server := httptest.NewServer(http.HandlerFunc(http.NotFound)) + defer server.Close() + + if _, err := runPProf(server.URL); err == nil { t.Fatalf("expected error for unknown file") } } @@ -130,7 +167,7 @@ func TestGetPProfRawBadURL(t *testing.T) { opts := Options{ BaseURL: "%-0", } - if _, err := GetRaw(opts); err == nil { + if _, err := GetRaw(opts, nil); err == nil { t.Error("expected bad BaseURL to fail") } } @@ -139,7 +176,7 @@ func TestGetPProfRawSuccess(t *testing.T) { opts := Options{ BinaryFile: "testdata/pprof.1.pb.gz", } - raw, err := GetRaw(opts) + raw, err := GetRaw(opts, nil) if err != nil { t.Fatalf("getPProfRaw failed: %v", err) }