Optimize WASM memory allocation

runtime/wasm/Cargo.toml

@@ -9,7 +9,7 @@ futures = "0.1.21"
 hex = "0.3.2"
 graph = { path = "../../graph" }
 tiny-keccak = "1.4.2"
-wasmi = "0.4"
+wasmi = "0.5"
 pwasm-utils = "0.6.1"
 bs58 = "0.2.2"
 graph-runtime-derive = { path = "../derive" }

runtime/wasm/src/host.rs

@@ -4,8 +4,9 @@ use semver::{Version, VersionReq};
 use tiny_keccak::keccak256;
 
 use std::collections::HashMap;
+use std::str::FromStr;
 use std::thread;
-use std::time::Instant;
+use std::time::{Duration, Instant};
 
 use super::MappingContext;
 use crate::module::{ValidModule, WasmiModule, WasmiModuleConfig};
@@ -19,6 +20,8 @@ use graph::prelude::{
 use graph::util;
 use web3::types::{Log, Transaction};
 
+pub(crate) const TIMEOUT_ENV_VAR: &str = "GRAPH_MAPPING_HANDLER_TIMEOUT";
+
 pub struct RuntimeHostConfig {
     subgraph_id: SubgraphDeploymentId,
     mapping: Mapping,
@@ -246,6 +249,10 @@ impl RuntimeHost {
                 ethereum_adapter: ethereum_adapter.clone(),
                 link_resolver: link_resolver.clone(),
                 store: store.clone(),
+                handler_timeout: std::env::var(TIMEOUT_ENV_VAR)
+                    .ok()
+                    .and_then(|s| u64::from_str(&s).ok())
+                    .map(Duration::from_secs),
             };
             let valid_module = ValidModule::new(&module_logger, wasmi_config, task_sender)
                 .expect("Failed to validate module");

runtime/wasm/src/host_exports.rs

@@ -17,8 +17,6 @@ use web3::types::H160;
 
 use crate::module::WasmiModule;
 
-pub(crate) const TIMEOUT_ENV_VAR: &str = "GRAPH_MAPPING_HANDLER_TIMEOUT";
-
 pub(crate) trait ExportError: fmt::Debug + fmt::Display + Send + Sync + 'static {}
 
 impl<E> ExportError for E where E: fmt::Debug + fmt::Display + Send + Sync + 'static {}
@@ -49,6 +47,7 @@ pub(crate) struct HostExports<E, L, S, U> {
     link_resolver: Arc<L>,
     store: Arc<S>,
     task_sink: U,
+    handler_timeout: Option<Duration>,
 }
 
 impl<E, L, S, U> HostExports<E, L, S, U>
@@ -72,6 +71,7 @@ where
         link_resolver: Arc<L>,
         store: Arc<S>,
         task_sink: U,
+        handler_timeout: Option<Duration>,
     ) -> Self {
         HostExports {
             subgraph_id,
@@ -83,6 +83,7 @@ where
             link_resolver,
             store,
             task_sink,
+            handler_timeout,
         }
     }
 
@@ -532,11 +533,7 @@ where
         &self,
         start_time: Instant,
     ) -> Result<(), HostExportError<impl ExportError>> {
-        let mapping_handler_timeout = std::env::var(TIMEOUT_ENV_VAR)
-            .ok()
-            .and_then(|s| u64::from_str(&s).ok())
-            .map(Duration::from_secs);
-        if let Some(timeout) = mapping_handler_timeout {
+        if let Some(timeout) = self.handler_timeout {
             if start_time.elapsed() > timeout {
                 return Err(HostExportError(format!("Mapping handler timed out")));
             }

runtime/wasm/src/lib.rs

@@ -1,14 +1,3 @@
-extern crate bs58;
-extern crate ethabi;
-extern crate futures;
-extern crate graph;
-extern crate graph_runtime_derive;
-extern crate hex;
-extern crate pwasm_utils;
-extern crate semver;
-extern crate tiny_keccak;
-extern crate wasmi;
-
 mod asc_abi;
 mod host;
 mod module;

runtime/wasm/src/module/mod.rs

@@ -1,6 +1,7 @@
+use std::convert::TryFrom;
 use std::fmt;
 use std::ops::Deref;
-use std::time::Instant;
+use std::time::{Duration, Instant};
 
 use semver::Version;
 use wasmi::{
@@ -86,6 +87,7 @@ pub struct WasmiModuleConfig<T, L, S> {
     pub ethereum_adapter: Arc<T>,
     pub link_resolver: Arc<L>,
     pub store: Arc<S>,
+    pub handler_timeout: Option<Duration>,
 }
 
 /// A pre-processed and valid WASM module, ready to be started as a WasmiModule.
@@ -161,6 +163,7 @@ where
             config.link_resolver.clone(),
             config.store.clone(),
             task_sink,
+            config.handler_timeout,
         );
 
         Ok(ValidModule {
@@ -187,6 +190,12 @@ pub(crate) struct WasmiModule<T, L, S, U> {
     // True if `run_start` has not yet been called on the module.
     // This is used to prevent mutating store state in start.
     running_start: bool,
+
+    // First free byte in the heap.
+    heap_start_ptr: u32,
+
+    // Number of free bytes starting from `heap_start_ptr`.
+    heap_free_size: u32,
 }
 
 impl<T, L, S, U> WasmiModule<T, L, S, U>
@@ -235,6 +244,10 @@ where
             valid_module: valid_module.clone(),
             start_time: Instant::now(),
             running_start: true,
+
+            // `heap_start_ptr` will be set on the first call to `raw_new`.
+            heap_free_size: 0,
+            heap_start_ptr: 0,
         };
 
         this.module = module
@@ -407,22 +420,37 @@ where
         + 'static,
 {
     fn raw_new(&mut self, bytes: &[u8]) -> Result<u32, Error> {
-        let address = self
-            .module
-            .clone()
-            .invoke_export(
-                "memory.allocate",
-                &[RuntimeValue::I32(bytes.len() as i32)],
-                self,
-            )
-            .expect("Failed to invoke memory allocation function")
-            .expect("Function did not return a value")
-            .try_into::<u32>()
-            .expect("Function did not return u32");
+        // We request large chunks from the AssemblyScript allocator and manage them ourselves.
+        // This assumes the arena allocator is being used in AS.
+
+        static MIN_HEAP_SIZE_INCREMENT: u32 = 10_000;
+
+        let size = u32::try_from(bytes.len()).unwrap();
+        if size > self.heap_free_size {
+            let need = size - self.heap_free_size;
+            let allocate = need.max(MIN_HEAP_SIZE_INCREMENT);
+            let allocated_ptr = self
+                .module
+                .clone()
+                .invoke_export("memory.allocate", &[RuntimeValue::from(allocate)], self)
+                .expect("Failed to invoke memory allocation function")
+                .expect("Function did not return a value")
+                .try_into::<u32>()
+                .expect("Function did not return u32");
+            self.heap_free_size += allocate;
+
+            // On the first call, initialze `self.heap_start_ptr`.
+            if self.heap_start_ptr == 0 {
+                self.heap_start_ptr = allocated_ptr;
+            }
+        };
 
-        self.memory.set(address, bytes)?;
+        let ptr = self.heap_start_ptr;
+        self.memory.set(ptr, bytes)?;
+        self.heap_start_ptr += size;
+        self.heap_free_size -= size;
 
-        Ok(address)
+        Ok(ptr)
     }
 
     fn get(&self, offset: u32, size: u32) -> Result<Vec<u8>, Error> {

runtime/wasm/src/module/test.rs

@@ -172,6 +172,10 @@ fn test_valid_module(
                 ethereum_adapter: mock_ethereum_adapter,
                 link_resolver: Arc::new(ipfs_api::IpfsClient::default().into()),
                 store: Arc::new(FakeStore),
+                handler_timeout: std::env::var(crate::host::TIMEOUT_ENV_VAR)
+                    .ok()
+                    .and_then(|s| u64::from_str(&s).ok())
+                    .map(Duration::from_secs),
             },
             task_sender,
         )

runtime/wasm/src/module/test/abi.rs

@@ -3,7 +3,7 @@ use super::*;
 #[test]
 fn unbounded_loop() {
     // Set handler timeout to 3 seconds.
-    env::set_var(host_exports::TIMEOUT_ENV_VAR, "3");
+    env::set_var(crate::host::TIMEOUT_ENV_VAR, "3");
     let valid_module = test_valid_module(mock_data_source("wasm_test/non_terminating.wasm"));
     let mut module = WasmiModule::from_valid_module_with_ctx(valid_module, mock_context()).unwrap();
     module.start_time = Instant::now();