add a command for taking a heap dump

agent/agent.go

@@ -204,6 +204,10 @@ func formatBytes(val uint64) string {
 	return fmt.Sprintf("%d bytes", val)
 }
 
+type fileConn interface {
+	File() (f *os.File, err error)
+}
+
 func handle(conn io.ReadWriter, msg []byte) error {
 	switch msg[0] {
 	case signal.StackTrace:
@@ -290,6 +294,19 @@ func handle(conn io.ReadWriter, msg []byte) error {
 			return err
 		}
 		fmt.Fprintf(conn, "New GC percent set to %v. Previous value was %v.\n", perc, debug.SetGCPercent(int(perc)))
+	case signal.WriteHeapDump:
+		fc, ok := conn.(fileConn)
+		if !ok {
+			return fmt.Errorf("unable to get file descriptor for connection: %T", conn)
+		}
+		f, err := fc.File()
+		if err != nil {
+			return err
+		}
+		defer f.Close()
+
+		runtime.GC()
+		debug.WriteHeapDump(f.Fd())
 	}
 	return nil
 }

internal/cmd/shared.go

@@ -70,6 +70,11 @@ func AgentCommands() []*cobra.Command {
 			short: "Reads the CPU profile and launches \"go tool pprof\".",
 			fn:    pprofCPU,
 		},
+		{
+			name:  "heap-dump",
+			short: "writes the results of debug.WriteHeapDump to a file.",
+			fn:    heapDump,
+		},
 		{
 			name:  "version",
 			short: "Prints the Go version used to build the program.",
@@ -249,6 +254,27 @@ func pprof(addr net.TCPAddr, p byte, prefix string) error {
 	return cmd.Run()
 }
 
+func heapDump(addr net.TCPAddr, params []string) error {
+	tmpDumpFile, err := os.CreateTemp("", "heap_dump")
+	if err != nil {
+		return err
+	}
+
+	out, err := cmd(addr, signal.WriteHeapDump)
+	if err != nil {
+		return err
+	}
+	if len(out) == 0 {
+		return errors.New("failed to read the dump")
+	}
+	if err := os.WriteFile(tmpDumpFile.Name(), out, 0); err != nil {
+		return err
+	}
+	fmt.Printf("Heap dump saved to: %s\n", tmpDumpFile.Name())
+
+	return nil
+}
+
 func stats(addr net.TCPAddr, _ []string) error {
 	return cmdWithPrint(addr, signal.Stats)
 }

signal/signal.go

@@ -35,4 +35,7 @@ const (
 
 	// SetGCPercent sets the garbage collection target percentage.
 	SetGCPercent = byte(0x10)
+
+	// WriteHeapDump writes the results of debug.WriteHeapDump to a file.
+	WriteHeapDump = byte(0x11)
 )