Merge pull request #2043 from OffchainLabs/metric-jit-wasm-memory

Add a metric for memory used by the replay binary WASM in JIT

arbitrator/jit/src/machine.rs

@@ -280,7 +280,7 @@ impl WasmEnv {
         Ok(env)
     }
 
-    pub fn send_results(&mut self, error: Option<String>) {
+    pub fn send_results(&mut self, error: Option<String>, memory_used: u64) {
         let writer = match &mut self.process.socket {
             Some((writer, _)) => writer,
             None => return,
@@ -307,6 +307,7 @@ impl WasmEnv {
         check!(socket::write_u64(writer, self.small_globals[1]));
         check!(socket::write_bytes32(writer, &self.large_globals[0]));
         check!(socket::write_bytes32(writer, &self.large_globals[1]));
+        check!(socket::write_u64(writer, memory_used));
         check!(writer.flush());
     }
 }

arbitrator/jit/src/main.rs

@@ -114,8 +114,9 @@ fn main() {
         true => None,
         false => Some(message),
     };
+    let memory_used = memory.size().0 as u64 * 65_536;
 
-    env.send_results(error);
+    env.send_results(error, memory_used);
 }
 
 // require a usize be at least 32 bits wide

arbitrator/jit/src/syscall.rs

@@ -337,7 +337,7 @@ pub fn js_value_call(mut env: WasmEnvMut, sp: u32) -> MaybeEscape {
 
     let value = match (object, method_name.as_slice()) {
         (Ref(GO_ID), b"_makeFuncWrapper") => {
-            let arg = match args.get(0) {
+            let arg = match args.first() {
                 Some(arg) => arg,
                 None => fail!(
                     "Go trying to call Go._makeFuncWrapper with bad args {:?}",
@@ -415,7 +415,7 @@ pub fn js_value_call(mut env: WasmEnvMut, sp: u32) -> MaybeEscape {
         (Ref(CRYPTO_ID), b"getRandomValues") => {
             let name = "crypto.getRandomValues";
 
-            let id = match args.get(0) {
+            let id = match args.first() {
                 Some(Ref(x)) => x,
                 _ => fail!("Go trying to call {name} with bad args {:?}", args),
             };
@@ -456,7 +456,7 @@ pub fn js_value_new(mut env: WasmEnvMut, sp: u32) {
     let args_len = sp.read_u64(2);
     let args = sp.read_value_slice(args_ptr, args_len);
     match class {
-        UINT8_ARRAY_ID => match args.get(0) {
+        UINT8_ARRAY_ID => match args.first() {
             Some(JsValue::Number(size)) => {
                 let id = pool.insert(DynamicObject::Uint8Array(vec![0; *size as usize]));
                 sp.write_u64(4, GoValue::Object(id).encode());

arbitrator/prover/src/machine.rs

@@ -362,7 +362,7 @@ impl Module {
             bin.memories.len() <= 1,
             "Multiple memories are not supported"
         );
-        if let Some(limits) = bin.memories.get(0) {
+        if let Some(limits) = bin.memories.first() {
             let page_size = Memory::PAGE_SIZE;
             let initial = limits.initial; // validate() checks this is less than max::u32
             let allowed = u32::MAX as u64 / Memory::PAGE_SIZE - 1; // we require the size remain *below* 2^32

arbitrator/wasm-libraries/go-stub/src/lib.rs

@@ -218,7 +218,7 @@ pub unsafe extern "C" fn go__syscall_js_valueNew(sp: GoStack) {
     let args_len = sp.read_u64(2);
     let args = read_value_slice(args_ptr, args_len);
     if class == UINT8_ARRAY_ID {
-        if let Some(InterpValue::Number(size)) = args.get(0) {
+        if let Some(InterpValue::Number(size)) = args.first() {
             let id = DynamicObjectPool::singleton()
                 .insert(DynamicObject::Uint8Array(vec![0; *size as usize]));
             sp.write_u64(4, GoValue::Object(id).encode());
@@ -321,7 +321,7 @@ unsafe fn value_call_impl(sp: &mut GoStack) -> Result<GoValue, String> {
     let args_len = sp.read_u64(4);
     let args = read_value_slice(args_ptr, args_len);
     if object == InterpValue::Ref(GO_ID) && &method_name == b"_makeFuncWrapper" {
-        let id = args.get(0).ok_or_else(|| {
+        let id = args.first().ok_or_else(|| {
             format!(
                 "Go attempting to call Go._makeFuncWrapper with bad args {:?}",
                 args,
@@ -405,7 +405,7 @@ unsafe fn value_call_impl(sp: &mut GoStack) -> Result<GoValue, String> {
             ))
         }
     } else if object == InterpValue::Ref(CRYPTO_ID) && &method_name == b"getRandomValues" {
-        let id = match args.get(0) {
+        let id = match args.first() {
             Some(InterpValue::Ref(x)) => *x,
             _ => {
                 return Err(format!(

validator/server_jit/jit_machine.go

@@ -16,17 +16,21 @@ import (
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/log"
+	"github.com/ethereum/go-ethereum/metrics"
 	"github.com/offchainlabs/nitro/util/arbmath"
 	"github.com/offchainlabs/nitro/validator"
 )
 
+var jitWasmMemoryUsage = metrics.NewRegisteredHistogram("jit/wasm/memoryusage", nil, metrics.NewBoundedHistogramSample())
+
 type JitMachine struct {
-	binary  string
-	process *exec.Cmd
-	stdin   io.WriteCloser
+	binary               string
+	process              *exec.Cmd
+	stdin                io.WriteCloser
+	wasmMemoryUsageLimit int
 }
 
-func createJitMachine(jitBinary string, binaryPath string, cranelift bool, moduleRoot common.Hash, fatalErrChan chan error) (*JitMachine, error) {
+func createJitMachine(jitBinary string, binaryPath string, cranelift bool, wasmMemoryUsageLimit int, moduleRoot common.Hash, fatalErrChan chan error) (*JitMachine, error) {
 	invocation := []string{"--binary", binaryPath, "--forks"}
 	if cranelift {
 		invocation = append(invocation, "--cranelift")
@@ -45,9 +49,10 @@ func createJitMachine(jitBinary string, binaryPath string, cranelift bool, modul
 	}()
 
 	machine := &JitMachine{
-		binary:  binaryPath,
-		process: process,
-		stdin:   stdin,
+		binary:               binaryPath,
+		process:              process,
+		stdin:                stdin,
+		wasmMemoryUsageLimit: wasmMemoryUsageLimit,
 	}
 	return machine, nil
 }
@@ -258,8 +263,18 @@ func (machine *JitMachine) prove(
 			if state.BlockHash, err = readHash(); err != nil {
 				return state, err
 			}
-			state.SendRoot, err = readHash()
-			return state, err
+			if state.SendRoot, err = readHash(); err != nil {
+				return state, err
+			}
+			memoryUsed, err := readUint64()
+			if err != nil {
+				return state, fmt.Errorf("failed to read memory usage from Jit machine: %w", err)
+			}
+			if memoryUsed > uint64(machine.wasmMemoryUsageLimit) {
+				log.Warn("memory used by jit wasm exceeds the wasm memory usage limit", "limit", machine.wasmMemoryUsageLimit, "memoryUsed", memoryUsed)
+			}
+			jitWasmMemoryUsage.Update(int64(memoryUsed))
+			return state, nil
 		default:
 			message := "inter-process communication failure"
 			log.Error("Jit Machine Failure", "message", message)

validator/server_jit/machine_loader.go

@@ -13,13 +13,15 @@ import (
 )
 
 type JitMachineConfig struct {
-	ProverBinPath string
-	JitCranelift  bool
+	ProverBinPath        string
+	JitCranelift         bool
+	WasmMemoryUsageLimit int
 }
 
 var DefaultJitMachineConfig = JitMachineConfig{
-	JitCranelift:  true,
-	ProverBinPath: "replay.wasm",
+	JitCranelift:         true,
+	ProverBinPath:        "replay.wasm",
+	WasmMemoryUsageLimit: 4294967296,
 }
 
 func getJitPath() (string, error) {
@@ -57,7 +59,7 @@ func NewJitMachineLoader(config *JitMachineConfig, locator *server_common.Machin
 	}
 	createMachineThreadFunc := func(ctx context.Context, moduleRoot common.Hash) (*JitMachine, error) {
 		binPath := filepath.Join(locator.GetMachinePath(moduleRoot), config.ProverBinPath)
-		return createJitMachine(jitPath, binPath, config.JitCranelift, moduleRoot, fatalErrChan)
+		return createJitMachine(jitPath, binPath, config.JitCranelift, config.WasmMemoryUsageLimit, moduleRoot, fatalErrChan)
 	}
 	return &JitMachineLoader{
 		MachineLoader: *server_common.NewMachineLoader[JitMachine](locator, createMachineThreadFunc),

validator/server_jit/spawner.go

@@ -18,18 +18,23 @@ import (
 type JitSpawnerConfig struct {
 	Workers   int  `koanf:"workers" reload:"hot"`
 	Cranelift bool `koanf:"cranelift"`
+
+	// TODO: change WasmMemoryUsageLimit to a string and use resourcemanager.ParseMemLimit
+	WasmMemoryUsageLimit int `koanf:"wasm-memory-usage-limit"`
 }
 
 type JitSpawnerConfigFecher func() *JitSpawnerConfig
 
 var DefaultJitSpawnerConfig = JitSpawnerConfig{
-	Workers:   0,
-	Cranelift: true,
+	Workers:              0,
+	Cranelift:            true,
+	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")
 }
 
 type JitSpawner struct {
@@ -44,6 +49,7 @@ func NewJitSpawner(locator *server_common.MachineLocator, config JitSpawnerConfi
 	// TODO - preload machines
 	machineConfig := DefaultJitMachineConfig
 	machineConfig.JitCranelift = config().Cranelift
+	machineConfig.WasmMemoryUsageLimit = config().WasmMemoryUsageLimit
 	loader, err := NewJitMachineLoader(&machineConfig, locator, fatalErrChan)
 	if err != nil {
 		return nil, err