chore: Enforce profiling.Prepare to enhance symbolization selection

cmd/wzprof/main.go

@@ -59,8 +59,29 @@ func (prog *program) run(ctx context.Context) error {
 
 	p := wzprof.ProfilingFor(wasmCode)
 
-	cpu := p.CPUProfiler(wzprof.HostTime(prog.hostTime))
-	mem := p.MemoryProfiler(wzprof.InuseMemory(prog.inuseMemory))
+	runtime := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfig().
+		WithDebugInfoEnabled(true).
+		WithCustomSections(true))
+
+	stdout.Printf("compiling wasm module %s", prog.filePath)
+	compiledModule, err := runtime.CompileModule(ctx, wasmCode)
+	if err != nil {
+		return fmt.Errorf("compiling wasm module: %w", err)
+	}
+	err = p.Prepare(compiledModule)
+	if err != nil {
+		return fmt.Errorf("preparing wasm module: %w", err)
+	}
+
+	cpu, err := p.CPUProfiler(wzprof.HostTime(prog.hostTime))
+	if err != nil {
+		return fmt.Errorf("creating cpu profiler: %w", err)
+	}
+
+	mem, err := p.MemoryProfiler(wzprof.InuseMemory(prog.inuseMemory))
+	if err != nil {
+		return fmt.Errorf("creating memory profiler: %w", err)
+	}
 
 	var listeners []experimental.FunctionListenerFactory
 	if prog.cpuProfile != "" || prog.pprofAddr != "" {
@@ -78,25 +99,6 @@ func (prog *program) run(ctx context.Context) error {
 		}
 	}
 
-	ctx = context.WithValue(ctx,
-		experimental.FunctionListenerFactoryKey{},
-		experimental.MultiFunctionListenerFactory(listeners...),
-	)
-
-	runtime := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfig().
-		WithDebugInfoEnabled(true).
-		WithCustomSections(true))
-
-	stdout.Printf("compiling wasm module %s", prog.filePath)
-	compiledModule, err := runtime.CompileModule(ctx, wasmCode)
-	if err != nil {
-		return fmt.Errorf("compiling wasm module: %w", err)
-	}
-	err = p.Prepare(compiledModule)
-	if err != nil {
-		return fmt.Errorf("preparing wasm module: %w", err)
-	}
-
 	if prog.pprofAddr != "" {
 		u := &url.URL{Scheme: "http", Host: prog.pprofAddr, Path: "/debug/pprof"}
 		stdout.Printf("starting prrof http sever at %s", u)

cpu_test.go

@@ -10,22 +10,31 @@ import (
 )
 
 func BenchmarkCPUProfilerOn(b *testing.B) {
-	p := ProfilingFor(nil).CPUProfiler()
+	p, err := ProfilingFor(nil).CPUProfiler()
+	if err != nil {
+		b.Fatal(err)
+	}
 	p.StartProfile()
 	benchmarkFunctionListener(b, p)
 }
 
 func BenchmarkCPUProfilerOff(b *testing.B) {
-	p := ProfilingFor(nil).CPUProfiler()
+	p, err := ProfilingFor(nil).CPUProfiler()
+	if err != nil {
+		b.Fatal(err)
+	}
 	benchmarkFunctionListener(b, p)
 }
 
 func TestCPUProfilerTime(t *testing.T) {
 	currentTime := int64(0)
 
-	p := ProfilingFor(nil).CPUProfiler(
+	p, err := ProfilingFor(nil).CPUProfiler(
 		TimeFunc(func() int64 { return currentTime }),
 	)
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	module := wazerotest.NewModule(nil,
 		wazerotest.NewFunction(func(context.Context, api.Module) {}),

mem_test.go

@@ -5,6 +5,9 @@ import (
 )
 
 func BenchmarkMemoryProfiler(b *testing.B) {
-	p := ProfilingFor(nil).MemoryProfiler()
+	p, err := ProfilingFor(nil).MemoryProfiler()
+	if err != nil {
+		b.Fatal(err)
+	}
 	benchmarkFunctionListener(b, p)
 }

wzprof.go

@@ -27,7 +27,8 @@ type Profiling struct {
 	symbols           symbolizer
 	stackIterator     func(mod api.Module, def api.FunctionDefinition, wasmsi experimental.StackIterator) experimental.StackIterator
 
-	lang language
+	lang          language
+	prepareCalled bool // Flag to indicate if Prepare has been called
 }
 
 type language int8
@@ -95,21 +96,10 @@ func ProfilingFor(wasm []byte) *Profiling {
 	return r
 }
 
-// CPUProfiler constructs a new instance of CPUProfiler using the given time
-// function to record the CPU time consumed.
-func (p *Profiling) CPUProfiler(options ...CPUProfilerOption) *CPUProfiler {
-	return newCPUProfiler(p, options...)
-}
-
-// MemoryProfiler constructs a new instance of MemoryProfiler using the given
-// time function to record the profile execution time.
-func (p *Profiling) MemoryProfiler(options ...MemoryProfilerOption) *MemoryProfiler {
-	return newMemoryProfiler(p, options...)
-}
-
 // Prepare selects the most appropriate analysis functions for the guest
 // code in the provided module.
 func (p *Profiling) Prepare(mod wazero.CompiledModule) error {
+
 	switch p.lang {
 	case golang:
 		s, err := preparePclntabSymbolizer(p.wasm, mod)
@@ -147,9 +137,31 @@ func (p *Profiling) Prepare(mod wazero.CompiledModule) error {
 		}
 		p.symbols = buildDwarfSymbolizer(dwarf)
 	}
+
+	// Set the flag to true if Prepare succeeds
+	p.prepareCalled = true
+
 	return nil
 }
 
+// CPUProfiler constructs a new instance of CPUProfiler using the given time
+// function to record the CPU time consumed.
+func (p *Profiling) CPUProfiler(options ...CPUProfilerOption) (*CPUProfiler, error) {
+	if !p.prepareCalled {
+		return nil, fmt.Errorf("Prepare must be called on the Profiling instance before creating a CPUProfiler")
+	}
+	return newCPUProfiler(p, options...), nil
+}
+
+// MemoryProfiler constructs a new instance of MemoryProfiler using the given
+// time function to record the profile execution time.
+func (p *Profiling) MemoryProfiler(options ...MemoryProfilerOption) (*MemoryProfiler, error) {
+	if !p.prepareCalled {
+		return nil, fmt.Errorf("Prepare must be called on the Profiling instance before creating a MemoryProfiler")
+	}
+	return newMemoryProfiler(p, options...), nil
+}
+
 // profilingListener wraps a FunctionListener to adapt its stack iterator to the
 // appropriate implementation according to the module support.
 type profilingListener struct {