Add timeout support in wazero run cli

.github/workflows/integration.yaml

@@ -112,8 +112,7 @@ jobs:
       # adding filter argument to the "Build Stdlib test binary" step.
       # e.g. --test-filter "Dir.Iterator but dir is deleted during iteration"
       - name: Run the test binary with wazero CLI
-        timeout-minutes: 10  # Prevent crashes from timing out
-        run: go run ./cmd/wazero run -mount=:/ test.wasm
+        run: go run ./cmd/wazero run -mount=:/ -timeout=10m test.wasm
 
   build_tinygo_test_binary:
     name: Build TinyGo test binary

cmd/wazero/testdata/infinite_loop.wat

@@ -0,0 +1,7 @@
+(module $infinite_loop
+  (func $main (export "_start")
+    (loop
+      br 0
+    )
+  )
+)

cmd/wazero/wazero.go

@@ -10,6 +10,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"time"
 
 	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
@@ -92,11 +93,11 @@ func doCompile(args []string, stdErr io.Writer, exit func(code int)) {
 
 	c := wazero.NewRuntimeConfig()
 	if cache := maybeUseCacheDir(cacheDir, stdErr, exit); cache != nil {
-		c.WithCompilationCache(cache)
+		c = c.WithCompilationCache(cache)
 	}
 
 	ctx := context.Background()
-	rt := wazero.NewRuntime(ctx)
+	rt := wazero.NewRuntimeWithConfig(ctx, c)
 	defer rt.Close(ctx)
 
 	if _, err = rt.CompileModule(ctx, wasm); err != nil {
@@ -133,6 +134,14 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer, exit func(cod
 			"This may be specified multiple times. When <wasm path> is unset, <path> is used. "+
 			"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. "+
+			"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\". "+
+			"If the duration is 0, the timeout is disabled. The default is disabled.")
+
 	var hostlogging logScopesFlag
 	flags.Var(&hostlogging, "hostlogging",
 		"a comma-separated list of host function scopes to log to stderr. "+
@@ -195,7 +204,17 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer, exit func(cod
 		rtc = wazero.NewRuntimeConfig()
 	}
 	if cache := maybeUseCacheDir(cacheDir, stdErr, exit); cache != nil {
-		rtc.WithCompilationCache(cache)
+		rtc = rtc.WithCompilationCache(cache)
+	}
+	if timeout > 0 {
+		newCtx, cancel := context.WithTimeout(ctx, timeout)
+		ctx = newCtx
+		defer cancel()
+		rtc = rtc.WithCloseOnContextDone(true)
+	} else if timeout < 0 {
+		fmt.Fprintf(stdErr, "timeout duration may not be negative, %v given\n", timeout)
+		printRunUsage(stdErr, flags)
+		exit(1)
 	}
 
 	rt := wazero.NewRuntimeWithConfig(ctx, rtc)
@@ -231,11 +250,17 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer, exit func(cod
 	} else if needsGo {
 		gojs.MustInstantiate(ctx, rt)
 		err = gojs.Run(ctx, rt, code, conf)
+	} else {
+		_, err = rt.InstantiateModule(ctx, code, conf)
 	}
 
 	if err != nil {
 		if exitErr, ok := err.(*sys.ExitError); ok {
-			exit(int(exitErr.ExitCode()))
+			exitCode := exitErr.ExitCode()
+			if exitCode == sys.ExitCodeDeadlineExceeded {
+				fmt.Fprintf(stdErr, "error: %v (timeout %v)\n", exitErr, timeout)
+			}
+			exit(int(exitCode))
 		}
 		fmt.Fprintf(stdErr, "error instantiating wasm binary: %v\n", err)
 		exit(1)

cmd/wazero/wazero_test.go

@@ -20,8 +20,12 @@ import (
 	"github.com/tetratelabs/wazero/internal/platform"
 	"github.com/tetratelabs/wazero/internal/testing/require"
 	"github.com/tetratelabs/wazero/internal/version"
+	"github.com/tetratelabs/wazero/sys"
 )
 
+//go:embed testdata/infinite_loop.wasm
+var wasmInfiniteLoop []byte
+
 //go:embed testdata/wasi_arg.wasm
 var wasmWasiArg []byte
 
@@ -426,6 +430,24 @@ func TestRun(t *testing.T) {
 				require.True(t, len(entries) > 0)
 			},
 		},
+		{
+			name:             "timeout: a binary that exceeds the deadline should print an error",
+			wazeroOpts:       []string{"-timeout=1ms"},
+			wasm:             wasmInfiniteLoop,
+			expectedStderr:   "error: module \"\" closed with context deadline exceeded (timeout 1ms)\n",
+			expectedExitCode: int(sys.ExitCodeDeadlineExceeded),
+			test: func(t *testing.T) {
+				require.NoError(t, err)
+			},
+		},
+		{
+			name:       "timeout: a binary that ends before the deadline should not print a timeout error",
+			wazeroOpts: []string{"-timeout=10s"},
+			wasm:       wasmWasiRandomGet,
+			test: func(t *testing.T) {
+				require.NoError(t, err)
+			},
+		},
 	}
 
 	cryptoTest := test{
@@ -512,6 +534,10 @@ func TestRun_Errors(t *testing.T) {
 			message: "invalid cachedir",
 			args:    []string{"--cachedir", notWasmPath, wasmPath},
 		},
+		{
+			message: "timeout duration may not be negative",
+			args:    []string{"-timeout=-10s", wasmPath},
+		},
 	}
 
 	for _, tc := range tests {