Extend gops trace to allow duration parameter

README.md

@@ -226,8 +226,10 @@ $ gops pprof-heap (<pid>|<addr>)
 
 ##### Execution trace
 
-gops allows you to start the runtime tracer for 5 seconds and examine the results.
+gops allows you to start the runtime tracer for a specified duration and
+examine the results.
 
 ```sh
-$ gops trace (<pid>|<addr>)
+$ gops trace (<pid>|<addr>) 90s
+Tracing now, will take 1m30s...
 ```

agent/agent.go

@@ -10,6 +10,7 @@ import (
 	"bufio"
 	"context"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -269,10 +270,24 @@ func handle(conn io.ReadWriter, msg []byte) error {
 		_, err = bufio.NewReader(f).WriteTo(conn)
 		return err
 	case signal.Trace:
+		intd, err := binary.ReadVarint(bufio.NewReader(conn))
+		if err != nil {
+			return err
+		}
+		d := time.Duration(intd)
+		if intd < 0 {
+			return errors.New("invalid duration")
+		}
+		if intd == 0 {
+			// also default to 5 seconds if the client didn't send any duration to
+			// keep compatibility between older clients contacting a newer server.
+			d = 5 * time.Second
+		}
+
 		if err := trace.Start(conn); err != nil {
 			return err
 		}
-		time.Sleep(5 * time.Second)
+		time.Sleep(d)
 		trace.Stop()
 	case signal.SetGCPercent:
 		perc, err := binary.ReadVarint(bufio.NewReader(conn))

cmd.go

@@ -15,6 +15,7 @@ import (
 	"os/exec"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/google/gops/internal"
 	"github.com/google/gops/signal"
@@ -71,9 +72,25 @@ func pprofCPU(addr net.TCPAddr, _ []string) error {
 	return pprof(addr, signal.CPUProfile, "cpu")
 }
 
-func trace(addr net.TCPAddr, _ []string) error {
-	fmt.Println("Tracing now, will take 5 secs...")
-	out, err := cmd(addr, signal.Trace)
+func trace(addr net.TCPAddr, params []string) error {
+	// keep the original duration of 5 seconds by default
+	duration := 5 * time.Second
+	buf := make([]byte, binary.MaxVarintLen64)
+
+	if len(params) > 0 {
+		d, err := time.ParseDuration(params[0])
+		if err != nil {
+			return fmt.Errorf("failed to parse duration: %v", params[0])
+		}
+		if d <= 0 {
+			return fmt.Errorf("duration has to be positive: %v", d)
+		}
+		duration = d
+		binary.PutVarint(buf, int64(duration))
+	}
+
+	fmt.Printf("Tracing now, will take %v...\n", duration)
+	out, err := cmd(addr, signal.Trace, buf...)
 	if err != nil {
 		return err
 	}

main.go

@@ -35,7 +35,7 @@ Commands:
   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".
+  trace      Runs the runtime tracer 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".