expose host pprof on http endpoint

README.md

@@ -60,8 +60,7 @@ test programs complete quickly.
 wzprof -sample 1 -memprofile /tmp/profile ./testdata/c/simple.wasm
 ```
 ```
-wzprof -sample 1 -cpuprofile /tmp/profile \
-    ./testdata/rust/simple/target/wasm32-wasi/debug/simple.wasm
+wzprof -sample 1 -cpuprofile /tmp/profile ./testdata/c/crunch_numbers.wasm
 ```
 ```
 go tool pprof -http :4000 /tmp/profile

cmd/wzprof/main.go

@@ -6,11 +6,12 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"math"
 	"net/http"
-	_ "net/http/pprof"
 	"os"
 	"os/signal"
 	"path/filepath"
+	"runtime"
 	"runtime/pprof"
 	"strings"
 
@@ -87,13 +88,11 @@ func (prog *program) run(ctx context.Context) error {
 	}
 
 	if prog.pprofAddr != "" {
-		pprof := http.NewServeMux()
-		pprof.Handle("/debug/pprof/profile", cpu.NewHandler(prog.sampleRate, symbols))
-		pprof.Handle("/debug/pprof/heap", mem.NewHandler(prog.sampleRate, symbols))
-		pprof.Handle("/wzprof", http.DefaultServeMux)
+		server := http.NewServeMux()
+		server.Handle("/debug/pprof/", wzprof.Index(prog.sampleRate, symbols, cpu, mem))
 
 		go func() {
-			if err := http.ListenAndServe(prog.pprofAddr, pprof); err != nil {
+			if err := http.ListenAndServe(prog.pprofAddr, server); err != nil {
 				log.Println(err)
 			}
 		}()
@@ -108,25 +107,33 @@ func (prog *program) run(ctx context.Context) error {
 			startCPUProfile(f)
 			defer stopCPUProfile()
 		}
+
 		if prog.memProfile != "" {
 			f, err := os.Create(prog.memProfile)
 			if err != nil {
 				return err
 			}
 			defer writeHeapProfile(f)
 		}
-	} else {
-		if prog.cpuProfile != "" {
-			cpu.StartProfile()
-			defer func() {
-				writeProfile(prog.cpuProfile, cpu.StopProfile(prog.sampleRate, symbols))
-			}()
-		}
-		if prog.memProfile != "" {
-			defer func() {
-				writeProfile(prog.memProfile, mem.NewProfile(prog.sampleRate, symbols))
-			}()
-		}
+	}
+
+	if prog.cpuProfile != "" {
+		cpu.StartProfile()
+		defer func() {
+			p := cpu.StopProfile(prog.sampleRate, symbols)
+			if !prog.hostProfile {
+				writeProfile(prog.cpuProfile, p)
+			}
+		}()
+	}
+
+	if prog.memProfile != "" {
+		defer func() {
+			p := mem.NewProfile(prog.sampleRate, symbols)
+			if !prog.hostProfile {
+				writeProfile(prog.memProfile, p)
+			}
+		}()
 	}
 
 	ctx, cancel := context.WithCancelCause(ctx)
@@ -199,6 +206,10 @@ func run(ctx context.Context) error {
 		return fmt.Errorf("usage: wzprof </path/to/app.wasm>")
 	}
 
+	rate := int(math.Ceil(1 / sampleRate))
+	runtime.SetBlockProfileRate(rate)
+	runtime.SetMutexProfileFraction(rate)
+
 	return (&program{
 		filePath:    args[0],
 		pprofAddr:   pprofAddr,

cpu.go

@@ -108,12 +108,34 @@ func (p *CPUProfiler) StopProfile(sampleRate float64, symbols Symbolizer) *profi
 		}
 	}
 
-	return buildProfile(sampleRate, symbols, samples, start, duration,
-		[]*profile.ValueType{
-			{Type: "samples", Unit: "count"},
-			{Type: "cpu", Unit: "nanoseconds"},
-		},
-	)
+	return buildProfile(sampleRate, symbols, samples, start, duration, p.SampleType())
+}
+
+// Name returns "profile" to match the name of the CPU profiler in pprof.
+func (p *CPUProfiler) Name() string {
+	return "profile"
+}
+
+// Desc returns a description copied from net/http/pprof.
+func (p *CPUProfiler) Desc() string {
+	return profileDescriptions[p.Name()]
+}
+
+// Count returns the number of execution stacks currently recorded in p.
+func (p *CPUProfiler) Count() int {
+	p.mutex.Lock()
+	n := len(p.counts)
+	p.mutex.Unlock()
+	return n
+}
+
+// SampleType returns the set of value types present in samples recorded by the
+// CPU profiler.
+func (p *CPUProfiler) SampleType() []*profile.ValueType {
+	return []*profile.ValueType{
+		{Type: "samples", Unit: "count"},
+		{Type: "cpu", Unit: "nanoseconds"},
+	}
 }
 
 // NewHandler returns a http handler allowing the profiler to be exposed on a
@@ -163,12 +185,12 @@ func (p *CPUProfiler) NewHandler(sampleRate float64, symbols Symbolizer) http.Ha
 // NewFunctionListener returns a function listener suited to record CPU timings
 // of calls to the function passed as argument.
 func (p *CPUProfiler) NewFunctionListener(def api.FunctionDefinition) experimental.FunctionListener {
-	return cpuListener{p}
+	return cpuProfiler{p}
 }
 
-type cpuListener struct{ *CPUProfiler }
+type cpuProfiler struct{ *CPUProfiler }
 
-func (p cpuListener) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) context.Context {
+func (p cpuProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) context.Context {
 	var frame cpuTimeFrame
 	p.mutex.Lock()
 
@@ -193,7 +215,7 @@ func (p cpuListener) Before(ctx context.Context, mod api.Module, def api.Functio
 	return ctx
 }
 
-func (p cpuListener) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, err error, results []uint64) {
+func (p cpuProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, err error, results []uint64) {
 	i := len(p.frames) - 1
 	f := p.frames[i]
 	p.frames = p.frames[:i]

go.mod

@@ -4,6 +4,6 @@ go 1.20
 
 require (
 	github.com/google/pprof v0.0.0-20230406165453-00490a63f317
-	github.com/tetratelabs/wazero v1.1.1-0.20230509090308-daad2fdb660d
+	github.com/tetratelabs/wazero v1.1.1-0.20230509151012-509827a23f68
 	golang.org/x/exp v0.0.0-20230425010034-47ecfdc1ba53
 )

go.sum

@@ -1,6 +1,6 @@
 github.com/google/pprof v0.0.0-20230406165453-00490a63f317 h1:hFhpt7CTmR3DX+b4R19ydQFtofxT0Sv3QsKNMVQYTMQ=
 github.com/google/pprof v0.0.0-20230406165453-00490a63f317/go.mod h1:79YE0hCXdHag9sBkw2o+N/YnZtTkXi0UT9Nnixa5eYk=
-github.com/tetratelabs/wazero v1.1.1-0.20230509090308-daad2fdb660d h1:C72IsWDI0IgtMA4gl+1viwZnZEsOmcQclYMKsj1aw2Q=
-github.com/tetratelabs/wazero v1.1.1-0.20230509090308-daad2fdb660d/go.mod h1:wYx2gNRg8/WihJfSDxA1TIL8H+GkfLYm+bIfbblu9VQ=
+github.com/tetratelabs/wazero v1.1.1-0.20230509151012-509827a23f68 h1:Yw5Aw4Ann50S7bTVsoe/8jmtvKnh9U4yv/97kjyYAIo=
+github.com/tetratelabs/wazero v1.1.1-0.20230509151012-509827a23f68/go.mod h1:wYx2gNRg8/WihJfSDxA1TIL8H+GkfLYm+bIfbblu9VQ=
 golang.org/x/exp v0.0.0-20230425010034-47ecfdc1ba53 h1:5llv2sWeaMSnA3w2kS57ouQQ4pudlXrR0dCgw51QK9o=
 golang.org/x/exp v0.0.0-20230425010034-47ecfdc1ba53/go.mod h1:V1LtkGg67GoY2N1AnLN78QLrzxkLyJw7RJb1gzOOz9w=

mem.go

@@ -67,6 +67,30 @@ func NewMemoryProfiler(opts ...MemoryProfilerOption) *MemoryProfiler {
 // NewProfile takes a snapshot of the current memory allocation state and builds
 // a profile representing the state of the program memory.
 func (p *MemoryProfiler) NewProfile(sampleRate float64, symbols Symbolizer) *profile.Profile {
+	return buildProfile(sampleRate, symbols, p.snapshot(), p.start, time.Since(p.start), p.SampleType())
+}
+
+// Name returns "allocs" to match the name of the memory profiler in pprof.
+func (p *MemoryProfiler) Name() string {
+	return "allocs"
+}
+
+// Desc returns a description copied from net/http/pprof.
+func (p *MemoryProfiler) Desc() string {
+	return profileDescriptions[p.Name()]
+}
+
+// Count returns the number of allocation stacks recorded in p.
+func (p *MemoryProfiler) Count() int {
+	p.mutex.Lock()
+	n := p.alloc.len()
+	p.mutex.Unlock()
+	return n
+}
+
+// SampleType returns the set of value types present in samples recorded by the
+// memory profiler.
+func (p *MemoryProfiler) SampleType() []*profile.ValueType {
 	sampleType := []*profile.ValueType{
 		{Type: "alloc_objects", Unit: "count"},
 		{Type: "alloc_space", Unit: "byte"},
@@ -82,7 +106,7 @@ func (p *MemoryProfiler) NewProfile(sampleRate float64, symbols Symbolizer) *pro
 		)
 	}
 
-	return buildProfile(sampleRate, symbols, p.snapshot(), p.start, time.Since(p.start), sampleType)
+	return sampleType
 }
 
 type memorySample struct {