Support pprof arg format in go-torch

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)
 	}

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.

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 != "" {

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)
 	}