From 0a33cbc9342c204c8d16ad3a61d94a569674b5f9 Mon Sep 17 00:00:00 2001
From: Glib Smaga <code@gsmaga.com>
Date: Tue, 29 Mar 2022 11:11:34 +0200
Subject: [PATCH] Extend `gops trace` to allow duration parameter

`gops trace :port 10m` will now run a trace for 10 minutes,
as opposed to always doing 5 seconds.

For backwards compatibility both server and the client
default to 5s if nothing is specified.
---
 README.md      |  6 ++++--
 agent/agent.go | 17 ++++++++++++++++-
 cmd.go         | 23 ++++++++++++++++++++---
 main.go        |  2 +-
 4 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 95abf296..7accb5ad 100644
--- a/README.md
+++ b/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...
 ```
diff --git a/agent/agent.go b/agent/agent.go
index 2af03e5a..2181bb6c 100644
--- a/agent/agent.go
+++ b/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))
diff --git a/cmd.go b/cmd.go
index ba63506c..a4dd8538 100644
--- a/cmd.go
+++ b/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
 	}
diff --git a/main.go b/main.go
index a85b0752..c5ed20b2 100644
--- a/main.go
+++ b/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".