runtime: Use interior mutability in AscHeapCtx

Author: lutter

runtime/wasm/src/module/context.rs

@@ -48,12 +48,8 @@ impl WasmInstanceContext<'_> {
         self.inner.data_mut()
     }
 
-    pub fn asc_heap_ref(&self) -> &AscHeapCtx {
-        self.as_ref().asc_heap_ref()
-    }
-
-    pub fn asc_heap_mut(&mut self) -> &mut AscHeapCtx {
-        self.as_mut().asc_heap_mut()
+    pub fn asc_heap(&self) -> &Arc<AscHeapCtx> {
+        self.as_ref().asc_heap()
     }
 
     pub fn suspend_timeout(&mut self) {
@@ -96,7 +92,7 @@ pub struct WasmInstanceData {
 
     // This option is needed to break the cyclic dependency between, instance, store, and context.
     // during execution it should always be populated.
-    asc_heap: Option<AscHeapCtx>,
+    asc_heap: Option<Arc<AscHeapCtx>>,
 }
 
 impl WasmInstanceData {
@@ -117,15 +113,12 @@ impl WasmInstanceData {
         }
     }
 
-    pub fn set_asc_heap(&mut self, asc_heap: AscHeapCtx) {
+    pub fn set_asc_heap(&mut self, asc_heap: Arc<AscHeapCtx>) {
         self.asc_heap = Some(asc_heap);
     }
 
-    pub fn asc_heap_ref(&self) -> &AscHeapCtx {
-        self.asc_heap.as_ref().unwrap()
-    }
-    pub fn asc_heap_mut(&mut self) -> &mut AscHeapCtx {
-        self.asc_heap.as_mut().unwrap()
+    pub fn asc_heap(&self) -> &Arc<AscHeapCtx> {
+        self.asc_heap.as_ref().expect("asc_heap not set")
     }
 
     pub fn take_state(mut self) -> BlockState {

runtime/wasm/src/module/mod.rs

@@ -1,5 +1,6 @@
 use std::convert::TryFrom;
 use std::mem::MaybeUninit;
+use std::sync::Mutex;
 
 use anyhow::anyhow;
 use anyhow::Error;
@@ -134,6 +135,19 @@ fn is_trap_deterministic(trap: &Error) -> bool {
     }
 }
 
+struct Arena {
+    // First free byte in the current arena. Set on the first call to `raw_new`.
+    start: i32,
+    // Number of free bytes starting from `arena_start_ptr`.
+    size: i32,
+}
+
+impl Arena {
+    fn new() -> Self {
+        Self { start: 0, size: 0 }
+    }
+}
+
 #[derive(Copy, Clone)]
 pub struct ExperimentalFeatures {
     pub allow_non_deterministic_ipfs: bool,
@@ -154,19 +168,15 @@ pub struct AscHeapCtx {
     // is zeroed when initialized or grown.
     memory: Memory,
 
-    // First free byte in the current arena. Set on the first call to `raw_new`.
-    arena_start_ptr: i32,
-
-    // Number of free bytes starting from `arena_start_ptr`.
-    arena_free_size: i32,
+    arena: Mutex<Arena>,
 }
 
 impl AscHeapCtx {
     pub(crate) fn new(
         instance: &wasmtime::Instance,
         ctx: &mut WasmInstanceContext<'_>,
         api_version: Version,
-    ) -> anyhow::Result<AscHeapCtx> {
+    ) -> anyhow::Result<Arc<AscHeapCtx>> {
         // Provide access to the WASM runtime linear memory
         let memory = instance
             .get_memory(ctx.as_context_mut(), "memory")
@@ -194,14 +204,33 @@ impl AscHeapCtx {
             ),
         };
 
-        Ok(AscHeapCtx {
+        Ok(Arc::new(AscHeapCtx {
             memory_allocate,
             memory,
-            arena_start_ptr: 0,
-            arena_free_size: 0,
+            arena: Mutex::new(Arena::new()),
             api_version,
             id_of_type,
-        })
+        }))
+    }
+
+    fn arena_start_ptr(&self) -> i32 {
+        self.arena.lock().unwrap().start
+    }
+
+    fn arena_free_size(&self) -> i32 {
+        self.arena.lock().unwrap().size
+    }
+
+    fn set_arena(&self, start_ptr: i32, size: i32) {
+        let mut arena = self.arena.lock().unwrap();
+        arena.start = start_ptr;
+        arena.size = size;
+    }
+
+    fn allocated(&self, size: i32) {
+        let mut arena = self.arena.lock().unwrap();
+        arena.start += size;
+        arena.size -= size;
     }
 }
 
@@ -229,21 +258,20 @@ impl AscHeap for WasmInstanceContext<'_> {
         static MIN_ARENA_SIZE: i32 = 10_000;
 
         let size = i32::try_from(bytes.len()).unwrap();
-        if size > self.asc_heap_ref().arena_free_size {
+        if size > self.asc_heap().arena_free_size() {
             // Allocate a new arena. Any free space left in the previous arena is left unused. This
             // causes at most half of memory to be wasted, which is acceptable.
-            let arena_size = size.max(MIN_ARENA_SIZE);
+            let mut arena_size = size.max(MIN_ARENA_SIZE);
 
             // Unwrap: This may panic if more memory needs to be requested from the OS and that
             // fails. This error is not deterministic since it depends on the operating conditions
             // of the node.
-            let memory_allocate = self.asc_heap_ref().memory_allocate.clone();
-            self.asc_heap_mut().arena_start_ptr = memory_allocate
+            let memory_allocate = &self.asc_heap().cheap_clone().memory_allocate;
+            let mut start_ptr = memory_allocate
                 .call(self.as_context_mut(), arena_size)
                 .unwrap();
-            self.asc_heap_mut().arena_free_size = arena_size;
 
-            match &self.asc_heap_ref().api_version {
+            match &self.asc_heap().api_version {
                 version if *version <= Version::new(0, 0, 4) => {}
                 _ => {
                     // This arithmetic is done because when you call AssemblyScripts's `__alloc`
@@ -252,27 +280,27 @@ impl AscHeap for WasmInstanceContext<'_> {
                     // `mmInfo` has size of 4, and everything allocated on AssemblyScript memory
                     // should have alignment of 16, this means we need to do a 12 offset on these
                     // big chunks of untyped allocation.
-                    self.asc_heap_mut().arena_start_ptr += 12;
-                    self.asc_heap_mut().arena_free_size -= 12;
+                    start_ptr += 12;
+                    arena_size -= 12;
                 }
             };
+            self.asc_heap().set_arena(start_ptr, arena_size);
         };
 
-        let ptr = self.asc_heap_ref().arena_start_ptr as usize;
+        let ptr = self.asc_heap().arena_start_ptr() as usize;
 
         // Unwrap: We have just allocated enough space for `bytes`.
-        let memory = self.asc_heap_ref().memory;
+        let memory = self.asc_heap().memory;
         memory.write(self.as_context_mut(), ptr, bytes).unwrap();
-        self.asc_heap_mut().arena_start_ptr += size;
-        self.asc_heap_mut().arena_free_size -= size;
+        self.asc_heap().allocated(size);
 
         Ok(ptr as u32)
     }
 
     fn read_u32(&self, offset: u32, gas: &GasCounter) -> Result<u32, DeterministicHostError> {
         gas.consume_host_fn_with_metrics(Gas::new(GAS_COST_LOAD as u64 * 4), "read_u32")?;
         let mut bytes = [0; 4];
-        self.asc_heap_ref()
+        self.asc_heap()
             .memory
             .read(self, offset as usize, &mut bytes)
             .map_err(|_| {
@@ -302,7 +330,7 @@ impl AscHeap for WasmInstanceContext<'_> {
 
         // TODO: Do we still need this? Can we use read directly?
         let src = self
-            .asc_heap_ref()
+            .asc_heap()
             .memory
             .data(self)
             .get(offset..)
@@ -317,11 +345,11 @@ impl AscHeap for WasmInstanceContext<'_> {
     }
 
     fn api_version(&self) -> Version {
-        self.asc_heap_ref().api_version.clone()
+        self.asc_heap().api_version.clone()
     }
 
     fn asc_type_id(&mut self, type_id_index: IndexForAscTypeId) -> Result<u32, HostExportError> {
-        let func = self.asc_heap_ref().id_of_type.clone().unwrap();
+        let func = self.asc_heap().id_of_type.clone().unwrap();
 
         // Unwrap ok because it's only called on correct apiVersion, look for AscPtr::generate_header
         func.call(self.as_context_mut(), type_id_index as u32)