feat(agent): support stack trace in debug=1 format with pprof labels

agent/agent.go

@@ -208,6 +208,8 @@ func handle(conn io.ReadWriter, msg []byte) error {
 	switch msg[0] {
 	case signal.StackTrace:
 		return pprof.Lookup("goroutine").WriteTo(conn, 2)
+	case signal.StackTraceDebug1:
+		return pprof.Lookup("goroutine").WriteTo(conn, 1)
 	case signal.GC:
 		runtime.GC()
 		_, err := conn.Write([]byte("ok"))

agent/agent_test.go

@@ -5,8 +5,12 @@
 package agent
 
 import (
+	"io"
+	"net"
 	"os"
 	"testing"
+
+	"github.com/google/gops/signal"
 )
 
 func TestListen(t *testing.T) {
@@ -87,3 +91,29 @@ func TestFormatBytes(t *testing.T) {
 		}
 	}
 }
+
+func TestStackTraceDebug1(t *testing.T) {
+	err := Listen(Options{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer Close()
+
+	conn, err := net.Dial("tcp", listener.Addr().String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+
+	if _, err := conn.Write([]byte{signal.StackTraceDebug1}); err != nil {
+		t.Fatal(err)
+	}
+
+	buf, err := io.ReadAll(conn)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(buf) == 0 {
+		t.Fatal("expected non-empty stack trace for debug=1")
+	}
+}

internal/cmd/shared.go

@@ -26,14 +26,20 @@ import (
 // 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.
+var res []*cobra.Command
+
+var debugLevel int
+
 func AgentCommands() []*cobra.Command {
-	var res []*cobra.Command
 
 	var cmds = []legacyCommand{
 		{
 			name:  "stack",
 			short: "Prints the stack trace.",
 			fn:    stackTrace,
+			setup: func(cmd *cobra.Command) {
+				cmd.Flags().IntVarP(&debugLevel, "debug", "d", 2, "debug level (1 or 2)")
+			},
 		},
 		{
 			name:  "gc",
@@ -79,7 +85,7 @@ func AgentCommands() []*cobra.Command {
 
 	for _, c := range cmds {
 		c := c
-		res = append(res, &cobra.Command{
+		cmd := &cobra.Command{
 			Use:   fmt.Sprintf("%s <pid|addr>", c.name),
 			Short: c.short,
 
@@ -110,7 +116,11 @@ func AgentCommands() []*cobra.Command {
 			// errors get double printed otherwise
 			SilenceUsage:  true,
 			SilenceErrors: true,
-		})
+		}
+		if c.setup != nil {
+			c.setup(cmd)
+		}
+		res = append(res, cmd)
 	}
 
 	return res
@@ -120,6 +130,7 @@ type legacyCommand struct {
 	name  string
 	short string
 	fn    func(addr net.TCPAddr, params []string) error
+	setup func(cmd *cobra.Command)
 }
 
 func setGC(addr net.TCPAddr, params []string) error {
@@ -144,7 +155,11 @@ func setGC(addr net.TCPAddr, params []string) error {
 }
 
 func stackTrace(addr net.TCPAddr, _ []string) error {
-	return cmdWithPrint(addr, signal.StackTrace)
+	sig := signal.StackTrace
+	if debugLevel == 1 {
+		sig = signal.StackTraceDebug1
+	}
+	return cmdWithPrint(addr, sig)
 }
 
 func gc(addr net.TCPAddr, _ []string) error {

signal/signal.go

@@ -35,4 +35,7 @@ const (
 
 	// SetGCPercent sets the garbage collection target percentage.
 	SetGCPercent = byte(0x10)
+
+	// StackTraceDebug1 represents a command to print stack trace with debug=1 (includes pprof labels).
+	StackTraceDebug1 = byte(0x11)
 )