From b42ec37d0b39c2cfce584d8e09d9dacd01514d9f Mon Sep 17 00:00:00 2001 From: Achille Date: Sat, 13 May 2023 23:23:57 -0700 Subject: [PATCH] wazero: allow cpu and memory profiling when compiling or running modules (#1462) Signed-off-by: Achille Roussel --- cmd/wazero/wazero.go | 168 +++++++++++++++++++++++++++----------- cmd/wazero/wazero_test.go | 66 +++++++++++---- 2 files changed, 168 insertions(+), 66 deletions(-) diff --git a/cmd/wazero/wazero.go b/cmd/wazero/wazero.go index 4845cbf1..2ef8c8f1 100644 --- a/cmd/wazero/wazero.go +++ b/cmd/wazero/wazero.go @@ -9,6 +9,7 @@ import ( "io" "os" "path/filepath" + "runtime/pprof" "strings" "time" @@ -24,51 +25,59 @@ import ( ) func main() { - doMain(os.Stdout, os.Stderr, os.Exit) + os.Exit(doMain(os.Stdout, os.Stderr)) } // doMain is separated out for the purpose of unit testing. -func doMain(stdOut io.Writer, stdErr logging.Writer, exit func(code int)) { +func doMain(stdOut io.Writer, stdErr logging.Writer) int { flag.CommandLine.SetOutput(stdErr) var help bool - flag.BoolVar(&help, "h", false, "print usage") + flag.BoolVar(&help, "h", false, "Prints usage.") flag.Parse() if help || flag.NArg() == 0 { printUsage(stdErr) - exit(0) + return 0 } if flag.NArg() < 1 { fmt.Fprintln(stdErr, "missing path to wasm file") printUsage(stdErr) - exit(1) + return 1 } subCmd := flag.Arg(0) switch subCmd { case "compile": - doCompile(flag.Args()[1:], stdErr, exit) + return doCompile(flag.Args()[1:], stdErr) case "run": - doRun(flag.Args()[1:], stdOut, stdErr, exit) + return doRun(flag.Args()[1:], stdOut, stdErr) case "version": fmt.Fprintln(stdOut, version.GetWazeroVersion()) - exit(0) + return 0 default: fmt.Fprintln(stdErr, "invalid command") printUsage(stdErr) - exit(1) + return 1 } } -func doCompile(args []string, stdErr io.Writer, exit func(code int)) { +func doCompile(args []string, stdErr io.Writer) int { flags := flag.NewFlagSet("compile", flag.ExitOnError) flags.SetOutput(stdErr) var help bool - flags.BoolVar(&help, "h", false, "print usage") + flags.BoolVar(&help, "h", false, "Prints usage.") + + var cpuProfile string + flags.StringVar(&cpuProfile, "cpuprofile", "", + "Enables cpu profiling and writes the profile at the given path.") + + var memProfile string + flags.StringVar(&memProfile, "memprofile", "", + "Enables memory profiling and writes the profile at the given path.") cacheDir := cacheDirFlag(flags) @@ -76,24 +85,36 @@ func doCompile(args []string, stdErr io.Writer, exit func(code int)) { if help { printCompileUsage(stdErr, flags) - exit(0) + return 0 } if flags.NArg() < 1 { fmt.Fprintln(stdErr, "missing path to wasm file") printCompileUsage(stdErr, flags) - exit(1) + return 1 } + + if memProfile != "" { + defer writeHeapProfile(stdErr, memProfile) + } + + if cpuProfile != "" { + stopCPUProfile := startCPUProfile(stdErr, cpuProfile) + defer stopCPUProfile() + } + wasmPath := flags.Arg(0) wasm, err := os.ReadFile(wasmPath) if err != nil { fmt.Fprintf(stdErr, "error reading wasm binary: %v\n", err) - exit(1) + return 1 } c := wazero.NewRuntimeConfig() - if cache := maybeUseCacheDir(cacheDir, stdErr, exit); cache != nil { + if rc, cache := maybeUseCacheDir(cacheDir, stdErr); rc != 0 { + return rc + } else if cache != nil { c = c.WithCompilationCache(cache) } @@ -103,22 +124,22 @@ func doCompile(args []string, stdErr io.Writer, exit func(code int)) { if _, err = rt.CompileModule(ctx, wasm); err != nil { fmt.Fprintf(stdErr, "error compiling wasm binary: %v\n", err) - exit(1) + return 1 } else { - exit(0) + return 0 } } -func doRun(args []string, stdOut io.Writer, stdErr logging.Writer, exit func(code int)) { +func doRun(args []string, stdOut io.Writer, stdErr logging.Writer) int { flags := flag.NewFlagSet("run", flag.ExitOnError) flags.SetOutput(stdErr) var help bool - flags.BoolVar(&help, "h", false, "print usage") + flags.BoolVar(&help, "h", false, "Prints usage.") var useInterpreter bool flags.BoolVar(&useInterpreter, "interpreter", false, - "interprets WebAssembly modules instead of compiling them into native code.") + "Interprets WebAssembly modules instead of compiling them into native code.") var envs sliceFlag flags.Var(&envs, "env", "key=value pair of environment variable to expose to the binary. "+ @@ -126,19 +147,19 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer, exit func(cod var envInherit bool flags.BoolVar(&envInherit, "env-inherit", false, - "inherits any environment variables from the calling process. "+ + "Inherits any environment variables from the calling process. "+ "Variables specified with the flag are appended to the inherited list.") var mounts sliceFlag flags.Var(&mounts, "mount", - "filesystem path to expose to the binary in the form of [:][:ro]. "+ + "Filesystem path to expose to the binary in the form of [:][:ro]. "+ "This may be specified multiple times. When is unset, is used. "+ "For example, -mount=/:/ or c:\\:/ makes the entire host volume writeable by wasm. "+ "For read-only mounts, append the suffix ':ro'.") var timeout time.Duration flags.DurationVar(&timeout, "timeout", 0*time.Second, - "if a wasm binary runs longer than the given duration string, then exit abruptly. "+ + "If a wasm binary runs longer than the given duration string, then exit abruptly. "+ "The duration string is an unsigned sequence of decimal numbers, "+ "each with optional fraction and a unit suffix, such as \"300ms\", \"1.5h\" or \"2h45m\". "+ "Valid time units are \"ns\", \"us\" (or \"µs\"), \"ms\", \"s\", \"m\", \"h\". "+ @@ -146,25 +167,42 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer, exit func(cod var hostlogging logScopesFlag flags.Var(&hostlogging, "hostlogging", - "a comma-separated list of host function scopes to log to stderr. "+ + "A comma-separated list of host function scopes to log to stderr. "+ "This may be specified multiple times. Supported values: all,clock,filesystem,memory,proc,poll,random") + var cpuProfile string + flags.StringVar(&cpuProfile, "cpuprofile", "", + "Enables cpu profiling and writes the profile at the given path.") + + var memProfile string + flags.StringVar(&memProfile, "memprofile", "", + "Enables memory profiling and writes the profile at the given path.") + cacheDir := cacheDirFlag(flags) _ = flags.Parse(args) if help { printRunUsage(stdErr, flags) - exit(0) + return 0 } if flags.NArg() < 1 { fmt.Fprintln(stdErr, "missing path to wasm file") printRunUsage(stdErr, flags) - exit(1) + return 1 } - wasmPath := flags.Arg(0) + if memProfile != "" { + defer writeHeapProfile(stdErr, memProfile) + } + + if cpuProfile != "" { + stopCPUProfile := startCPUProfile(stdErr, cpuProfile) + defer stopCPUProfile() + } + + wasmPath := flags.Arg(0) wasmArgs := flags.Args()[1:] if len(wasmArgs) > 1 { // Skip "--" if provided @@ -182,17 +220,20 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer, exit func(cod fields := strings.SplitN(e, "=", 2) if len(fields) != 2 { fmt.Fprintf(stdErr, "invalid environment variable: %s\n", e) - exit(1) + return 1 } env = append(env, fields[0], fields[1]) } - rootPath, fsConfig := validateMounts(mounts, stdErr, exit) + rc, rootPath, fsConfig := validateMounts(mounts, stdErr) + if rc != 0 { + return rc + } wasm, err := os.ReadFile(wasmPath) if err != nil { fmt.Fprintf(stdErr, "error reading wasm binary: %v\n", err) - exit(1) + return 1 } wasmExe := filepath.Base(wasmPath) @@ -206,7 +247,9 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer, exit func(cod ctx := maybeHostLogging(context.Background(), logging.LogScopes(hostlogging), stdErr) - if cache := maybeUseCacheDir(cacheDir, stdErr, exit); cache != nil { + if rc, cache := maybeUseCacheDir(cacheDir, stdErr); rc != 0 { + return rc + } else if cache != nil { rtc = rtc.WithCompilationCache(cache) } @@ -218,7 +261,7 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer, exit func(cod } else if timeout < 0 { fmt.Fprintf(stdErr, "timeout duration may not be negative, %v given\n", timeout) printRunUsage(stdErr, flags) - exit(1) + return 1 } rt := wazero.NewRuntimeWithConfig(ctx, rtc) @@ -243,7 +286,7 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer, exit func(cod code, err := rt.CompileModule(ctx, wasm) if err != nil { fmt.Fprintf(stdErr, "error compiling wasm binary: %v\n", err) - exit(1) + return 1 } switch detectImports(code.ImportedFunctions()) { @@ -286,22 +329,22 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer, exit func(cod if exitCode == sys.ExitCodeDeadlineExceeded { fmt.Fprintf(stdErr, "error: %v (timeout %v)\n", exitErr, timeout) } - exit(int(exitCode)) + return int(exitCode) } fmt.Fprintf(stdErr, "error instantiating wasm binary: %v\n", err) - exit(1) + return 1 } // We're done, _start was called as part of instantiating the module. - exit(0) + return 0 } -func validateMounts(mounts sliceFlag, stdErr logging.Writer, exit func(code int)) (rootPath string, config wazero.FSConfig) { +func validateMounts(mounts sliceFlag, stdErr logging.Writer) (rc int, rootPath string, config wazero.FSConfig) { config = wazero.NewFSConfig() for _, mount := range mounts { if len(mount) == 0 { fmt.Fprintln(stdErr, "invalid mount: empty string") - exit(1) + return 1, rootPath, config } readOnly := false @@ -327,14 +370,14 @@ func validateMounts(mounts sliceFlag, stdErr logging.Writer, exit func(code int) // Eagerly validate the mounts as we know they should be on the host. if abs, err := filepath.Abs(dir); err != nil { fmt.Fprintf(stdErr, "invalid mount: path %q invalid: %v\n", dir, err) - exit(1) + return 1, rootPath, config } else { dir = abs } if stat, err := os.Stat(dir); err != nil { fmt.Fprintf(stdErr, "invalid mount: path %q error: %v\n", dir, err) - exit(1) + return 1, rootPath, config } else if !stat.IsDir() { fmt.Fprintf(stdErr, "invalid mount: path %q is not a directory\n", dir) } @@ -349,7 +392,7 @@ func validateMounts(mounts sliceFlag, stdErr logging.Writer, exit func(code int) rootPath = dir } } - return + return 0, rootPath, config } const ( @@ -388,18 +431,16 @@ func cacheDirFlag(flags *flag.FlagSet) *string { "Contents are re-used for the same version of wazero.") } -func maybeUseCacheDir(cacheDir *string, stdErr io.Writer, exit func(code int)) (cache wazero.CompilationCache) { +func maybeUseCacheDir(cacheDir *string, stdErr io.Writer) (int, wazero.CompilationCache) { if dir := *cacheDir; dir != "" { - var err error - cache, err = wazero.NewCompilationCacheWithDir(dir) - if err != nil { + if cache, err := wazero.NewCompilationCacheWithDir(dir); err != nil { fmt.Fprintf(stdErr, "invalid cachedir: %v\n", err) - exit(1) + return 1, cache } else { - return + return 0, cache } } - return + return 0, nil } func printUsage(stdErr io.Writer) { @@ -431,6 +472,37 @@ func printRunUsage(stdErr io.Writer, flags *flag.FlagSet) { flags.PrintDefaults() } +func startCPUProfile(stdErr io.Writer, path string) (stopCPUProfile func()) { + f, err := os.Create(path) + if err != nil { + fmt.Fprintf(stdErr, "error creating cpu profile output: %v\n", err) + return func() {} + } + + if err := pprof.StartCPUProfile(f); err != nil { + f.Close() + fmt.Fprintf(stdErr, "error starting cpu profile: %v\n", err) + return func() {} + } + + return func() { + defer f.Close() + pprof.StopCPUProfile() + } +} + +func writeHeapProfile(stdErr io.Writer, path string) { + f, err := os.Create(path) + if err != nil { + fmt.Fprintf(stdErr, "error creating memory profile output: %v\n", err) + return + } + defer f.Close() + if err := pprof.WriteHeapProfile(f); err != nil { + fmt.Fprintf(stdErr, "error writing memory profile: %v\n", err) + } +} + type sliceFlag []string func (f *sliceFlag) String() string { diff --git a/cmd/wazero/wazero_test.go b/cmd/wazero/wazero_test.go index ef6f700e..5264130a 100644 --- a/cmd/wazero/wazero_test.go +++ b/cmd/wazero/wazero_test.go @@ -75,6 +75,9 @@ func TestCompile(t *testing.T) { existingDir2 := filepath.Join(tmpDir, "existing2") require.NoError(t, os.Mkdir(existingDir2, 0o700)) + cpuProfile := filepath.Join(t.TempDir(), "cpu.out") + memProfile := filepath.Join(t.TempDir(), "mem.out") + tests := []struct { name string wazeroOpts []string @@ -119,6 +122,20 @@ func TestCompile(t *testing.T) { require.True(t, len(entries) > 0) }, }, + { + name: "enable cpu profiling", + wazeroOpts: []string{"-cpuprofile=" + cpuProfile}, + test: func(t *testing.T) { + require.NoError(t, exist(cpuProfile)) + }, + }, + { + name: "enable memory profiling", + wazeroOpts: []string{"-memprofile=" + memProfile}, + test: func(t *testing.T) { + require.NoError(t, exist(memProfile)) + }, + }, } for _, tc := range tests { @@ -226,6 +243,9 @@ func TestRun(t *testing.T) { existingDir2 := filepath.Join(tmpDir, "existing2") require.NoError(t, os.Mkdir(existingDir2, 0o700)) + cpuProfile := filepath.Join(t.TempDir(), "cpu.out") + memProfile := filepath.Join(t.TempDir(), "mem.out") + type test struct { name string wazeroOpts []string @@ -482,6 +502,22 @@ func TestRun(t *testing.T) { require.NoError(t, err) }, }, + { + name: "enable cpu profiling", + wazeroOpts: []string{"-cpuprofile=" + cpuProfile}, + wasm: wasmWasiRandomGet, + test: func(t *testing.T) { + require.NoError(t, exist(cpuProfile)) + }, + }, + { + name: "enable memory profiling", + wazeroOpts: []string{"-memprofile=" + memProfile}, + wasm: wasmWasiRandomGet, + test: func(t *testing.T) { + require.NoError(t, exist(memProfile)) + }, + }, } cryptoTest := test{ @@ -762,25 +798,11 @@ func runMain(t *testing.T, workdir string, args []string) (int, string, string) os.Args = oldArgs }) os.Args = append([]string{"wazero"}, args...) + flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ContinueOnError) - var exitCode int - var stdout, stderr bytes.Buffer - var exited bool - func() { - defer func() { - if r := recover(); r != nil { - exited = true - } - }() - flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ContinueOnError) - doMain(&stdout, &stderr, func(code int) { - exitCode = code - panic(code) // to exit the func and set the exit status. - }) - }() - - require.True(t, exited) - + stdout := new(bytes.Buffer) + stderr := new(bytes.Buffer) + exitCode := doMain(stdout, stderr) return exitCode, stdout.String(), stderr.String() } @@ -806,3 +828,11 @@ func compileGoJS() (err error) { wasmCatGo, err = os.ReadFile(outPath) return } + +func exist(path string) error { + f, err := os.Open(path) + if err != nil { + return err + } + return f.Close() +}