diff --git a/agent/agent.go b/agent/agent.go index 9e605039..bbfef671 100644 --- a/agent/agent.go +++ b/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")) diff --git a/agent/agent_test.go b/agent/agent_test.go index e335a56c..e9eced84 100644 --- a/agent/agent_test.go +++ b/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") + } +} diff --git a/internal/cmd/shared.go b/internal/cmd/shared.go index fc11e6b2..6c8236e7 100644 --- a/internal/cmd/shared.go +++ b/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 ", 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 { diff --git a/signal/signal.go b/signal/signal.go index c70764a0..9f17f8ae 100644 --- a/signal/signal.go +++ b/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) )