vm: deduplicate code on host function tracing (#11503)

I'm exploring movement of the import code to their respective implementations in order to make the interfaces to the import code a little more flexible (and possibly less dependent on VMLogic struct.) Part of #11319
near · Jun 7, 2024 · 91c7726 · 91c7726
1 parent 337a81e
commit 91c7726
Showing 1 changed file with 32 additions and 38 deletions.
diff --git a/runtime/near-vm-runner/src/imports.rs b/runtime/near-vm-runner/src/imports.rs
@@ -299,7 +299,6 @@ imports! {
 
 #[cfg(all(feature = "wasmer0_vm", target_arch = "x86_64"))]
 pub(crate) mod wasmer {
-    use super::str_eq;
     use crate::logic::{VMLogic, VMLogicError};
     use std::ffi::c_void;
 
@@ -329,12 +328,10 @@ pub(crate) mod wasmer {
             ) => {
                 #[allow(unused_parens)]
                 fn $name( ctx: &mut wasmer_runtime::Ctx, $( $arg_name: $arg_type ),* ) -> Result<($( $returns ),*), VMLogicError> {
-                    const IS_GAS: bool = str_eq(stringify!($name), "gas") || str_eq(stringify!($name), "finite_wasm_gas");
-                    let _span = if IS_GAS {
-                        None
-                    } else {
-                        Some(tracing::trace_span!(target: "host-function", stringify!($name)).entered())
-                    };
+                    const TRACE: bool = $crate::imports::should_trace_host_function(stringify!($name));
+                    let _span = TRACE.then(|| {
+                        tracing::trace_span!(target: "vm::host_function", stringify!($name)).entered()
+                    });
                     let logic: &mut VMLogic<'_> = unsafe { &mut *(ctx.data as *mut VMLogic<'_>) };
                     logic.$func( $( $arg_name, )* )
                 }
@@ -356,10 +353,8 @@ pub(crate) mod wasmer {
 
 #[cfg(all(feature = "wasmer2_vm", target_arch = "x86_64"))]
 pub(crate) mod wasmer2 {
-    use std::sync::Arc;
-
-    use super::str_eq;
     use crate::logic::VMLogic;
+    use std::sync::Arc;
     use wasmer_engine::Engine;
     use wasmer_engine_universal::UniversalEngine;
     use wasmer_vm::{
@@ -427,15 +422,10 @@ pub(crate) mod wasmer2 {
                     extern "C" fn $name(env: *mut VMLogic<'_>, $( $arg_name: $arg_type ),* )
                     -> Ret {
                         let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
-                            const IS_GAS: bool = str_eq(stringify!($name), "gas") || str_eq(stringify!($name), "finite_wasm_gas");
-                            let _span = if IS_GAS {
-                                None
-                            } else {
-                                Some(tracing::trace_span!(
-                                    target: "host-function",
-                                    stringify!($name)
-                                ).entered())
-                            };
+                            const TRACE: bool = $crate::imports::should_trace_host_function(stringify!($name));
+                            let _span = TRACE.then(|| {
+                                tracing::trace_span!(target: "vm::host_function", stringify!($name)).entered()
+                            });
 
                             // SAFETY: This code should only be executable within `'vmlogic`
                             // lifetime and so it is safe to dereference the `env` pointer which is
@@ -510,14 +500,12 @@ pub(crate) mod wasmer2 {
 
 #[cfg(all(feature = "near_vm", target_arch = "x86_64"))]
 pub(crate) mod near_vm {
-    use std::sync::Arc;
-
-    use super::str_eq;
     use crate::logic::VMLogic;
     use near_vm_engine::universal::UniversalEngine;
     use near_vm_vm::{
         ExportFunction, ExportFunctionMetadata, Resolver, VMFunction, VMFunctionKind, VMMemory,
     };
+    use std::sync::Arc;
 
     pub(crate) struct NearVmImports<'engine, 'vmlogic, 'vmlogic_refs> {
         pub(crate) memory: VMMemory,
@@ -580,15 +568,10 @@ pub(crate) mod near_vm {
                     extern "C" fn $name(env: *mut VMLogic<'_>, $( $arg_name: $arg_type ),* )
                     -> Ret {
                         let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
-                            const IS_GAS: bool = str_eq(stringify!($name), "gas") || str_eq(stringify!($name), "finite_wasm_gas");
-                            let _span = if IS_GAS {
-                                None
-                            } else {
-                                Some(tracing::trace_span!(
-                                    target: "host-function",
-                                    stringify!($name)
-                                ).entered())
-                            };
+                            const TRACE: bool = $crate::imports::should_trace_host_function(stringify!($name));
+                            let _span = TRACE.then(|| {
+                                tracing::trace_span!(target: "vm::host_function", stringify!($name)).entered()
+                            });
 
                             // SAFETY: This code should only be executable within `'vmlogic`
                             // lifetime and so it is safe to dereference the `env` pointer which is
@@ -663,7 +646,6 @@ pub(crate) mod near_vm {
 
 #[cfg(feature = "wasmtime_vm")]
 pub(crate) mod wasmtime {
-    use super::str_eq;
     use crate::logic::{VMLogic, VMLogicError};
     use std::cell::UnsafeCell;
     use std::ffi::c_void;
@@ -710,12 +692,10 @@ pub(crate) mod wasmtime {
             ) => {
                 #[allow(unused_parens)]
                 fn $name(caller: wasmtime::Caller<'_, ()>, $( $arg_name: $arg_type ),* ) -> anyhow::Result<($( $returns ),*)> {
-                    const IS_GAS: bool = str_eq(stringify!($name), "gas") || str_eq(stringify!($name), "finite_wasm_gas");
-                    let _span = if IS_GAS {
-                        None
-                    } else {
-                        Some(tracing::trace_span!(target: "host-function", stringify!($name)).entered())
-                    };
+                    const TRACE: bool = $crate::imports::should_trace_host_function(stringify!($name));
+                    let _span = TRACE.then(|| {
+                        tracing::trace_span!(target: "vm::host_function", stringify!($name)).entered()
+                    });
                     // the below is bad. don't do this at home. it probably works thanks to the exact way the system is setup.
                     // Thanksfully, this doesn't run in production, and hopefully should be possible to remove before we even
                     // consider doing so.
@@ -744,6 +724,20 @@ pub(crate) mod wasmtime {
     }
 }
 
+#[cfg(any(
+    feature = "wasmer0_vm",
+    feature = "wasmer2_vm",
+    feature = "near_vm",
+    feature = "wasmtime_vm"
+))]
+pub(crate) const fn should_trace_host_function(host_function: &str) -> bool {
+    match host_function {
+        _ if str_eq(host_function, "gas") => false,
+        _ if str_eq(host_function, "finite_wasm_gas") => false,
+        _ => true,
+    }
+}
+
 /// Constant-time string equality, work-around for `"foo" == "bar"` not working
 /// in const context yet.
 #[cfg(any(