Merge pull request #2698 from renlulu/fix/validation-timeout

refactor: make timeout configurable for jit execution
OffchainLabs · Sep 24, 2024 · ab72b49 · ab72b49
2 parents e7d7198 + 3ea5870
commit ab72b49
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 9 deletions.
diff --git a/validator/server_jit/jit_machine.go b/validator/server_jit/jit_machine.go
@@ -30,9 +30,10 @@ type JitMachine struct {
 	process              *exec.Cmd
 	stdin                io.WriteCloser
 	wasmMemoryUsageLimit int
+	maxExecutionTime     time.Duration
 }
 
-func createJitMachine(jitBinary string, binaryPath string, cranelift bool, wasmMemoryUsageLimit int, moduleRoot common.Hash, fatalErrChan chan error) (*JitMachine, error) {
+func createJitMachine(jitBinary string, binaryPath string, cranelift bool, wasmMemoryUsageLimit int, maxExecutionTime time.Duration, moduleRoot common.Hash, fatalErrChan chan error) (*JitMachine, error) {
 	invocation := []string{"--binary", binaryPath, "--forks"}
 	if cranelift {
 		invocation = append(invocation, "--cranelift")
@@ -55,6 +56,7 @@ func createJitMachine(jitBinary string, binaryPath string, cranelift bool, wasmM
 		process:              process,
 		stdin:                stdin,
 		wasmMemoryUsageLimit: wasmMemoryUsageLimit,
+		maxExecutionTime:     maxExecutionTime,
 	}
 	return machine, nil
 }
@@ -73,7 +75,7 @@ func (machine *JitMachine) prove(
 	defer cancel() // ensure our cleanup functions run when we're done
 	state := validator.GoGlobalState{}
 
-	timeout := time.Now().Add(60 * time.Second)
+	timeout := time.Now().Add(machine.maxExecutionTime)
 	tcp, err := net.ListenTCP("tcp4", &net.TCPAddr{
 		IP: []byte{127, 0, 0, 1},
 	})

diff --git a/validator/server_jit/machine_loader.go b/validator/server_jit/machine_loader.go
@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"strings"
+	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/offchainlabs/nitro/validator/server_common"
@@ -52,14 +53,14 @@ type JitMachineLoader struct {
 	stopped bool
 }
 
-func NewJitMachineLoader(config *JitMachineConfig, locator *server_common.MachineLocator, fatalErrChan chan error) (*JitMachineLoader, error) {
+func NewJitMachineLoader(config *JitMachineConfig, locator *server_common.MachineLocator, maxExecutionTime time.Duration, fatalErrChan chan error) (*JitMachineLoader, error) {
 	jitPath, err := getJitPath()
 	if err != nil {
 		return nil, err
 	}
 	createMachineThreadFunc := func(ctx context.Context, moduleRoot common.Hash) (*JitMachine, error) {
 		binPath := filepath.Join(locator.GetMachinePath(moduleRoot), config.ProverBinPath)
-		return createJitMachine(jitPath, binPath, config.JitCranelift, config.WasmMemoryUsageLimit, moduleRoot, fatalErrChan)
+		return createJitMachine(jitPath, binPath, config.JitCranelift, config.WasmMemoryUsageLimit, maxExecutionTime, moduleRoot, fatalErrChan)
 	}
 	return &JitMachineLoader{
 		MachineLoader: *server_common.NewMachineLoader[JitMachine](locator, createMachineThreadFunc),

diff --git a/validator/server_jit/spawner.go b/validator/server_jit/spawner.go
@@ -3,10 +3,10 @@ package server_jit
 import (
 	"context"
 	"fmt"
+	flag "github.com/spf13/pflag"
 	"runtime"
 	"sync/atomic"
-
-	flag "github.com/spf13/pflag"
+	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/rawdb"
@@ -18,8 +18,9 @@ import (
 )
 
 type JitSpawnerConfig struct {
-	Workers   int  `koanf:"workers" reload:"hot"`
-	Cranelift bool `koanf:"cranelift"`
+	Workers          int           `koanf:"workers" reload:"hot"`
+	Cranelift        bool          `koanf:"cranelift"`
+	MaxExecutionTime time.Duration `koanf:"max-execution-time" reload:"hot"`
 
 	// TODO: change WasmMemoryUsageLimit to a string and use resourcemanager.ParseMemLimit
 	WasmMemoryUsageLimit int `koanf:"wasm-memory-usage-limit"`
@@ -30,13 +31,15 @@ type JitSpawnerConfigFecher func() *JitSpawnerConfig
 var DefaultJitSpawnerConfig = JitSpawnerConfig{
 	Workers:              0,
 	Cranelift:            true,
+	MaxExecutionTime:     time.Minute * 10,
 	WasmMemoryUsageLimit: 4294967296, // 2^32 WASM memeory limit
 }
 
 func JitSpawnerConfigAddOptions(prefix string, f *flag.FlagSet) {
 	f.Int(prefix+".workers", DefaultJitSpawnerConfig.Workers, "number of concurrent validation threads")
 	f.Bool(prefix+".cranelift", DefaultJitSpawnerConfig.Cranelift, "use Cranelift instead of LLVM when validating blocks using the jit-accelerated block validator")
 	f.Int(prefix+".wasm-memory-usage-limit", DefaultJitSpawnerConfig.WasmMemoryUsageLimit, "if memory used by a jit wasm exceeds this limit, a warning is logged")
+	f.Duration(prefix+".max-execution-time", DefaultJitSpawnerConfig.MaxExecutionTime, "if execution time used by a jit wasm exceeds this limit, a rpc error is returned")
 }
 
 type JitSpawner struct {
@@ -52,7 +55,8 @@ func NewJitSpawner(locator *server_common.MachineLocator, config JitSpawnerConfi
 	machineConfig := DefaultJitMachineConfig
 	machineConfig.JitCranelift = config().Cranelift
 	machineConfig.WasmMemoryUsageLimit = config().WasmMemoryUsageLimit
-	loader, err := NewJitMachineLoader(&machineConfig, locator, fatalErrChan)
+	maxExecutionTime := config().MaxExecutionTime
+	loader, err := NewJitMachineLoader(&machineConfig, locator, maxExecutionTime, fatalErrChan)
 	if err != nil {
 		return nil, err
 	}