[WIP] feat: pprof http path

pkg/proc/objects.go

@@ -125,16 +125,20 @@ type finalMarkParam struct {
 func (s *ObjRefScope) finalMark(idx *pprofIndex, hb *heapBits) {
 	var ptr Address
 	var size, count int64
+	var mem proc.MemoryReadWriter
 	for {
 		ptr = hb.nextPtr(true)
 		if ptr == 0 {
 			break
 		}
-		ptr, err := readUintRaw(s.mem, uint64(ptr), int64(s.bi.Arch.PtrSize()))
+		if mem == nil {
+			mem = cacheMemory(s.mem, uint64(ptr), int(hb.end.Sub(ptr)))
+		}
+		ptr, err := readUintRaw(mem, uint64(ptr), int64(s.bi.Arch.PtrSize()))
 		if err != nil {
 			continue
 		}
-		size_, count_ := s.markObject(Address(ptr), s.mem)
+		size_, count_ := s.markObject(Address(ptr), mem)
 		size += size_
 		count += count_
 	}

pkg/reexec/command_linux.go

@@ -0,0 +1,26 @@
+package reexec
+
+import (
+	"os/exec"
+	"syscall"
+)
+
+// Command returns an [*exec.Cmd] which has Path as current binary which,
+// on Linux, is set to the in-memory version (/proc/self/exe) of the current
+// binary, it is thus safe to delete or replace the on-disk binary (os.Args[0]).
+//
+// On Linux, the Pdeathsig of [*exec.Cmd.SysProcAttr] is set to SIGTERM.
+// This signal will be sent to the process when the OS thread which created
+// the process dies.
+//
+// It is the caller's responsibility to ensure that the creating thread is
+// not terminated prematurely. See https://go.dev/issue/27505 for more details.
+func Command(args ...string) *exec.Cmd {
+	return &exec.Cmd{
+		Path: Self(),
+		Args: args,
+		SysProcAttr: &syscall.SysProcAttr{
+			Pdeathsig: syscall.SIGTERM,
+		},
+	}
+}

pkg/reexec/command_other.go

@@ -0,0 +1,19 @@
+//go:build freebsd || darwin || windows
+
+package reexec
+
+import (
+	"os/exec"
+)
+
+// Command returns *exec.Cmd with its Path set to the path of the current
+// binary using the result of [Self]. For example if current binary is
+// "my-binary" at "/usr/bin/" (or "my-binary.exe" at "C:\" on Windows),
+// then cmd.Path is set to "/usr/bin/my-binary" and "C:\my-binary.exe"
+// respectively.
+func Command(args ...string) *exec.Cmd {
+	return &exec.Cmd{
+		Path: Self(),
+		Args: args,
+	}
+}

pkg/reexec/command_unsupported.go

@@ -0,0 +1,12 @@
+//go:build !linux && !windows && !freebsd && !darwin
+
+package reexec
+
+import (
+	"os/exec"
+)
+
+// Command is unsupported on operating systems apart from Linux, Windows, and Darwin.
+func Command(args ...string) *exec.Cmd {
+	return nil
+}

pkg/reexec/reexec.go

@@ -0,0 +1,64 @@
+// Package reexec facilitates the busybox style reexec of a binary.
+//
+// Handlers can be registered with a name and the argv 0 of the exec of
+// the binary will be used to find and execute custom init paths.
+//
+// It is used in dockerd to work around forking limitations when using Go.
+package reexec
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"runtime"
+)
+
+var registeredInitializers = make(map[string]func())
+
+// Register adds an initialization func under the specified name. It panics
+// if the given name is already registered.
+func Register(name string, initializer func()) {
+	if _, exists := registeredInitializers[name]; exists {
+		panic(fmt.Sprintf("reexec func already registered under name %q", name))
+	}
+
+	registeredInitializers[name] = initializer
+}
+
+// Init is called as the first part of the exec process and returns true if an
+// initialization function was called.
+func Init() bool {
+	if initializer, ok := registeredInitializers[os.Args[0]]; ok {
+		initializer()
+		return true
+	}
+	return false
+}
+
+// Self returns the path to the current process's binary. On Linux, it
+// returns "/proc/self/exe", which provides the in-memory version of the
+// current binary, whereas on other platforms it attempts to looks up the
+// absolute path for os.Args[0], or otherwise returns os.Args[0] as-is.
+func Self() string {
+	if runtime.GOOS == "linux" {
+		return "/proc/self/exe"
+	}
+	return naiveSelf()
+}
+
+func naiveSelf() string {
+	name := os.Args[0]
+	if filepath.Base(name) == name {
+		if lp, err := exec.LookPath(name); err == nil {
+			return lp
+		}
+	}
+	// handle conversion of relative paths to absolute
+	if absName, err := filepath.Abs(name); err == nil {
+		return absName
+	}
+	// if we couldn't get absolute name, return original
+	// (NOTE: Go only errors on Abs() if os.Getwd fails)
+	return name
+}

pkg/reexec/reexec_test.go

@@ -0,0 +1,67 @@
+package reexec
+
+import (
+	"os"
+	"os/exec"
+	"testing"
+)
+
+func init() {
+	Register("reexec", func() {
+		panic("Return Error")
+	})
+	Init()
+}
+
+func TestRegister(t *testing.T) {
+	defer func() {
+		if r := recover(); r != nil {
+			const expected = `reexec func already registered under name "reexec"`
+			if r != expected {
+				t.Errorf("got %q, want %q", r, expected)
+			}
+		}
+	}()
+	Register("reexec", func() {})
+}
+
+func TestCommand(t *testing.T) {
+	cmd := Command("reexec")
+	w, err := cmd.StdinPipe()
+	if err != nil {
+		t.Fatalf("Error on pipe creation: %v", err)
+	}
+	defer w.Close()
+
+	err = cmd.Start()
+	if err != nil {
+		t.Fatalf("Error on re-exec cmd: %v", err)
+	}
+	err = cmd.Wait()
+	const expected = "exit status 2"
+	if err == nil || err.Error() != expected {
+		t.Fatalf("got %v, want %v", err, expected)
+	}
+}
+
+func TestNaiveSelf(t *testing.T) {
+	if os.Getenv("TEST_CHECK") == "1" {
+		os.Exit(2)
+	}
+	cmd := exec.Command(naiveSelf(), "-test.run=TestNaiveSelf")
+	cmd.Env = append(os.Environ(), "TEST_CHECK=1")
+	err := cmd.Start()
+	if err != nil {
+		t.Fatalf("Unable to start command: %v", err)
+	}
+	err = cmd.Wait()
+	const expected = "exit status 2"
+	if err == nil || err.Error() != expected {
+		t.Fatalf("got %v, want %v", err, expected)
+	}
+
+	os.Args[0] = "mkdir"
+	if naiveSelf() == os.Args[0] {
+		t.Fatalf("Expected naiveSelf to resolve the location of mkdir")
+	}
+}

pprof/profiling.go

@@ -0,0 +1,45 @@
+package pprof
+
+import (
+	"log"
+	"net/http"
+	"os"
+
+	"github.com/cloudwego/goref/pkg/reexec"
+)
+
+const forkGoref = "fork_goref"
+
+func init() {
+	reexec.Register(forkGoref, func() {
+		println("hello world")
+		// attachCmd
+	})
+	if reexec.Init() {
+		os.Exit(0)
+	}
+}
+
+func init() {
+	http.HandleFunc("/debug/pprof/reference", Reference)
+}
+
+// Reference responds with the pprof-formatted heap reference profile.
+func Reference(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("X-Content-Type-Options", "nosniff")
+
+	// Set Content Type assuming StartCPUProfile will work,
+	// because if it does it starts writing.
+	w.Header().Set("Content-Type", "application/octet-stream")
+	w.Header().Set("Content-Disposition", `attachment; filename="profile"`)
+	cmd := reexec.Command(forkGoref)
+	cmd.Stdin = os.Stdin
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	if err := cmd.Start(); err != nil {
+		log.Panicf("failed to run goref command: %v", err)
+	}
+	if err := cmd.Wait(); err != nil {
+		log.Panicf("failed to wait goref command: %v", err)
+	}
+}