feat: RUN-681: New Wasm Instrumentation Part 3

dfinity · Aug 16, 2023 · 9d106b8 · 9d106b8
1 parent 72c08a5
commit 9d106b8
Show file tree

Hide file tree

Showing 11 changed files with 685 additions and 109 deletions.
diff --git a/rs/embedders/src/wasm_utils/instrumentation.rs b/rs/embedders/src/wasm_utils/instrumentation.rs
@@ -485,6 +485,7 @@ pub(super) fn instrument(
             special_indices,
             subnet_type,
             dirty_page_overhead,
+            metering_type,
         )
     }
 
@@ -563,6 +564,7 @@ fn replace_system_api_functions(
     special_indices: SpecialIndices,
     subnet_type: SubnetType,
     dirty_page_overhead: NumInstructions,
+    metering_type: MeteringType,
 ) {
     let api_indexes = calculate_api_indexes(module);
     let number_of_func_imports = module
@@ -574,9 +576,12 @@ fn replace_system_api_functions(
     // Collect a single map of all the function indexes that need to be
     // replaced.
     let mut func_index_replacements = BTreeMap::new();
-    for (api, (ty, body)) in
-        replacement_functions(special_indices, subnet_type, dirty_page_overhead)
-    {
+    for (api, (ty, body)) in replacement_functions(
+        special_indices,
+        subnet_type,
+        dirty_page_overhead,
+        metering_type,
+    ) {
         if let Some(old_index) = api_indexes.get(&api) {
             let type_idx = add_type(module, ty);
             let new_index = (number_of_func_imports + module.functions.len()) as u32;

diff --git a/rs/embedders/src/wasm_utils/system_api_replacements.rs b/rs/embedders/src/wasm_utils/system_api_replacements.rs
@@ -12,9 +12,10 @@
 //!
 
 use crate::{
-    wasm_utils::instrumentation::InjectedImports,
-    wasmtime_embedder::system_api_complexity::overhead, InternalErrorCode,
+    wasm_utils::instrumentation::InjectedImports, wasmtime_embedder::system_api_complexity,
+    InternalErrorCode,
 };
+use ic_config::embedders::MeteringType;
 use ic_interfaces::execution_environment::StableMemoryApi;
 use ic_registry_subnet_type::SubnetType;
 use ic_sys::PAGE_SIZE;
@@ -24,12 +25,15 @@ use wasmtime_environ::WASM_PAGE_SIZE;
 
 use super::{instrumentation::SpecialIndices, wasm_transform::Body, SystemApiFunc};
 
+use crate::wasmtime_embedder::system_api_complexity::system_api;
+
 const MAX_32_BIT_STABLE_MEMORY_IN_PAGES: i64 = 64 * 1024; // 4GiB
 
 pub(super) fn replacement_functions(
     special_indices: SpecialIndices,
     subnet_type: SubnetType,
     dirty_page_overhead: NumInstructions,
+    metering_type: MeteringType,
 ) -> Vec<(SystemApiFunc, (Type, Body<'static>))> {
     let count_clean_pages_fn_index = special_indices.count_clean_pages_fn.unwrap();
     let dirty_pages_counter_index = special_indices.dirty_pages_counter_ix.unwrap();
@@ -241,7 +245,11 @@ pub(super) fn replacement_functions(
                             },
                             I64ExtendI32U,
                             I64Const {
-                                value: overhead::STABLE_READ.get() as i64,
+                                value: system_api::complexity_overhead_native!(
+                                    STABLE_READ,
+                                    metering_type
+                                )
+                                .get() as i64,
                             },
                             I64Add,
                             Call {
@@ -502,7 +510,11 @@ pub(super) fn replacement_functions(
                                 }
                             },
                             I64Const {
-                                value: overhead::STABLE64_READ.get() as i64,
+                                value: system_api::complexity_overhead_native!(
+                                    STABLE64_READ,
+                                    metering_type
+                                )
+                                .get() as i64,
                             },
                             I64Add,
                             Call {
@@ -791,7 +803,11 @@ pub(super) fn replacement_functions(
                             },
                             I64ExtendI32U,
                             I64Const {
-                                value: overhead::STABLE_WRITE.get() as i64,
+                                value: system_api::complexity_overhead_native!(
+                                    STABLE_WRITE,
+                                    metering_type
+                                )
+                                .get() as i64,
                             },
                             I64Add,
                             Call {
@@ -1021,7 +1037,11 @@ pub(super) fn replacement_functions(
                                 }
                             },
                             I64Const {
-                                value: overhead::STABLE64_WRITE.get() as i64,
+                                value: system_api::complexity_overhead_native!(
+                                    STABLE64_WRITE,
+                                    metering_type
+                                )
+                                .get() as i64,
                             },
                             I64Add,
                             Call {

diff --git a/rs/embedders/src/wasmtime_embedder.rs b/rs/embedders/src/wasmtime_embedder.rs
@@ -343,6 +343,7 @@ impl WasmtimeEmbedder {
             self.config.feature_flags,
             self.config.stable_memory_dirty_page_limit,
             self.config.stable_memory_accessed_page_limit,
+            self.config.metering_type,
         );
 
         let instance = match linker.instantiate(&mut store, module) {