revert ctx cmd/main.go

Signed-off-by: Rajiv Singh <[email protected]>

cmd/wzprof/main.go

@@ -59,20 +59,6 @@ func (prog *program) run(ctx context.Context) error {
 
 	p := wzprof.ProfilingFor(wasmCode)
 
-	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 := p.CPUProfiler(wzprof.HostTime(prog.hostTime))
 	mem := p.MemoryProfiler(wzprof.InuseMemory(prog.inuseMemory))
 
@@ -92,6 +78,25 @@ 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)