From 1cf23c0fd14cf16de1206bae535ccaa3d6fcdf7f Mon Sep 17 00:00:00 2001
From: Glen Mailer <glen@geckoboard.com>
Date: Mon, 1 Dec 2025 17:02:29 +0000
Subject: [PATCH] add a command for taking a heap dump

---
 agent/agent.go         | 17 +++++++++++++++++
 internal/cmd/shared.go | 26 ++++++++++++++++++++++++++
 signal/signal.go       |  3 +++
 3 files changed, 46 insertions(+)

diff --git a/agent/agent.go b/agent/agent.go
index 9e605039..57c68b98 100644
--- a/agent/agent.go
+++ b/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
 }
diff --git a/internal/cmd/shared.go b/internal/cmd/shared.go
index fc11e6b2..5499f07b 100644
--- a/internal/cmd/shared.go
+++ b/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)
 }
diff --git a/signal/signal.go b/signal/signal.go
index c70764a0..89ccfd20 100644
--- a/signal/signal.go
+++ b/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)
 )