feat: add gas limit

poc/runtime/src/pvq.rs

@@ -44,7 +44,8 @@ pub mod extensions {
 
 pub fn execute_query(query: &[u8], input: &[u8]) -> PvqResult {
     let mut executor = ExtensionsExecutor::<extensions::Extensions, ()>::new(InvokeSource::RuntimeAPI);
-    executor.execute_method(query, input, 0)
+    let (result, _) = executor.execute_method(query, input, None);
+    result
 }
 
 pub fn metadata() -> Metadata {

pvq-executor/src/context.rs

@@ -0,0 +1,8 @@
+use polkavm::Linker;
+
+pub trait PvqExecutorContext {
+    type UserData;
+    type UserError;
+    fn register_host_functions(&mut self, linker: &mut Linker<Self::UserData, Self::UserError>);
+    fn data(&mut self) -> &mut Self::UserData;
+}

pvq-executor/src/error.rs

@@ -0,0 +1,50 @@
+use pvq_primitives::PvqError;
+#[derive(Debug)]
+pub enum PvqExecutorError<UserError> {
+    InvalidProgramFormat,
+    MemoryAccessError(polkavm::MemoryAccessError),
+    // Extract from the PVM CallError
+    Trap,
+    // Extract from the PVM CallError
+    NotEnoughGas,
+    // Usually a custom error type from the extension system definition
+    User(UserError),
+    // Other errors directly from the PVM
+    OtherPvmError(polkavm::Error),
+}
+
+impl<UserError> From<polkavm::MemoryAccessError> for PvqExecutorError<UserError> {
+    fn from(err: polkavm::MemoryAccessError) -> Self {
+        Self::MemoryAccessError(err)
+    }
+}
+
+impl<UserError> From<polkavm::Error> for PvqExecutorError<UserError> {
+    fn from(err: polkavm::Error) -> Self {
+        Self::OtherPvmError(err)
+    }
+}
+
+impl<UserError> From<polkavm::CallError<UserError>> for PvqExecutorError<UserError> {
+    fn from(err: polkavm::CallError<UserError>) -> Self {
+        match err {
+            polkavm::CallError::Trap => Self::Trap,
+            polkavm::CallError::NotEnoughGas => Self::NotEnoughGas,
+            polkavm::CallError::Error(e) => Self::OtherPvmError(e),
+            polkavm::CallError::User(e) => Self::User(e),
+        }
+    }
+}
+
+impl<UserError> From<PvqExecutorError<UserError>> for PvqError {
+    fn from(e: PvqExecutorError<UserError>) -> PvqError {
+        match e {
+            PvqExecutorError::InvalidProgramFormat => PvqError::InvalidPvqProgramFormat,
+            PvqExecutorError::MemoryAccessError(_) => PvqError::MemoryAccessError,
+            PvqExecutorError::Trap => PvqError::Trap,
+            PvqExecutorError::NotEnoughGas => PvqError::QueryExceedsWeightLimit,
+            PvqExecutorError::User(_) => PvqError::HostCallError,
+            PvqExecutorError::OtherPvmError(_) => PvqError::Other,
+        }
+    }
+}

pvq-executor/src/executor.rs

@@ -0,0 +1,92 @@
+use alloc::vec::Vec;
+use polkavm::{Config, Engine, Linker, Module, ModuleConfig, ProgramBlob};
+
+use crate::context::PvqExecutorContext;
+use crate::error::PvqExecutorError;
+
+type PvqExecutorResult<UserError> = Result<Vec<u8>, PvqExecutorError<UserError>>;
+type GasLimit = Option<i64>;
+
+pub struct PvqExecutor<Ctx: PvqExecutorContext> {
+    engine: Engine,
+    linker: Linker<Ctx::UserData, Ctx::UserError>,
+    context: Ctx,
+}
+
+impl<Ctx: PvqExecutorContext> PvqExecutor<Ctx> {
+    pub fn new(config: Config, mut context: Ctx) -> Self {
+        let engine = Engine::new(&config).unwrap();
+        let mut linker = Linker::<Ctx::UserData, Ctx::UserError>::new();
+        // Register user-defined host functions
+        context.register_host_functions(&mut linker);
+        Self {
+            engine,
+            linker,
+            context,
+        }
+    }
+
+    pub fn execute(
+        &mut self,
+        program: &[u8],
+        args: &[u8],
+        gas_limit: GasLimit,
+    ) -> (PvqExecutorResult<Ctx::UserError>, GasLimit) {
+        let blob = match ProgramBlob::parse(program.into()) {
+            Ok(blob) => blob,
+            Err(_) => return (Err(PvqExecutorError::InvalidProgramFormat), gas_limit),
+        };
+
+        // TODO: make this configurable
+        let mut module_config = ModuleConfig::new();
+        module_config.set_aux_data_size(args.len() as u32);
+        if gas_limit.is_some() {
+            module_config.set_gas_metering(Some(polkavm::GasMeteringKind::Sync));
+        }
+
+        let module = match Module::from_blob(&self.engine, &module_config, blob) {
+            Ok(module) => module,
+            Err(err) => return (Err(err.into()), gas_limit),
+        };
+
+        let instance_pre = match self.linker.instantiate_pre(&module) {
+            Ok(instance_pre) => instance_pre,
+            Err(err) => return (Err(err.into()), gas_limit),
+        };
+
+        let mut instance = match instance_pre.instantiate() {
+            Ok(instance) => instance,
+            Err(err) => return (Err(err.into()), gas_limit),
+        };
+
+        if let Some(gas_limit) = gas_limit {
+            instance.set_gas(gas_limit);
+        }
+
+        // From this point on, we include instance.gas() in the return value
+        let result = (|| {
+            instance.write_memory(module.memory_map().aux_data_address(), args)?;
+
+            tracing::info!("Calling entrypoint with args: {:?}", args);
+            let res = instance.call_typed_and_get_result::<u64, (u32, u32)>(
+                self.context.data(),
+                "pvq",
+                (module.memory_map().aux_data_address(), args.len() as u32),
+            )?;
+
+            let res_size = (res >> 32) as u32;
+            let res_ptr = (res & 0xffffffff) as u32;
+
+            let result = instance.read_memory(res_ptr, res_size)?;
+
+            tracing::info!("Result: {:?}", result);
+            Ok(result)
+        })();
+
+        if gas_limit.is_some() {
+            (result, Some(instance.gas()))
+        } else {
+            (result, None)
+        }
+    }
+}

pvq-executor/src/lib.rs

@@ -3,88 +3,12 @@
 extern crate alloc;
 
 pub use alloc::vec::Vec;
-use polkavm::ModuleConfig;
 pub use polkavm::{Caller, Config, Engine, Linker, Module, ProgramBlob};
 
-pub trait PvqExecutorContext {
-    type UserData;
-    type UserError;
-    fn register_host_functions(&mut self, linker: &mut Linker<Self::UserData, Self::UserError>);
-    fn data(&mut self) -> &mut Self::UserData;
-}
+mod context;
+mod error;
+mod executor;
 
-pub struct PvqExecutor<Ctx: PvqExecutorContext> {
-    engine: Engine,
-    linker: Linker<Ctx::UserData, Ctx::UserError>,
-    context: Ctx,
-}
-
-#[derive(Debug)]
-pub enum PvqExecutorError<UserError> {
-    MemoryAllocationError,
-    MemoryAccessError(polkavm::MemoryAccessError),
-    CallError(polkavm::CallError<UserError>),
-    OtherPVMError(polkavm::Error),
-}
-
-impl<UserError> From<polkavm::MemoryAccessError> for PvqExecutorError<UserError> {
-    fn from(err: polkavm::MemoryAccessError) -> Self {
-        Self::MemoryAccessError(err)
-    }
-}
-
-impl<UserError> From<polkavm::Error> for PvqExecutorError<UserError> {
-    fn from(err: polkavm::Error) -> Self {
-        Self::OtherPVMError(err)
-    }
-}
-
-impl<UserError> From<polkavm::CallError<UserError>> for PvqExecutorError<UserError> {
-    fn from(err: polkavm::CallError<UserError>) -> Self {
-        Self::CallError(err)
-    }
-}
-
-impl<Ctx: PvqExecutorContext> PvqExecutor<Ctx> {
-    pub fn new(config: Config, mut context: Ctx) -> Self {
-        let engine = Engine::new(&config).unwrap();
-        let mut linker = Linker::<Ctx::UserData, Ctx::UserError>::new();
-        // Register user-defined host functions
-        context.register_host_functions(&mut linker);
-        Self {
-            engine,
-            linker,
-            context,
-        }
-    }
-
-    pub fn execute(
-        &mut self,
-        program: &[u8],
-        args: &[u8],
-        _gas_limit: u64,
-    ) -> Result<Vec<u8>, PvqExecutorError<Ctx::UserError>> {
-        let blob = ProgramBlob::parse(program.into()).map_err(polkavm::Error::from)?;
-
-        // TODO: make this configurable
-        let mut module_config = ModuleConfig::new();
-        module_config.set_aux_data_size(args.len() as u32);
-
-        let module = Module::from_blob(&self.engine, &module_config, blob)?;
-        let instance_pre = self.linker.instantiate_pre(&module)?;
-        let mut instance = instance_pre.instantiate()?;
-
-        instance.write_memory(module.memory_map().aux_data_address(), args)?;
-        tracing::info!("Calling entrypoint with args: {:?}", args);
-        let res = instance.call_typed_and_get_result::<u64, (u32, u32)>(
-            self.context.data(),
-            "pvq",
-            (module.memory_map().aux_data_address(), args.len() as u32),
-        )?;
-        let res_size = (res >> 32) as u32;
-        let res_ptr = (res & 0xffffffff) as u32;
-        let result = instance.read_memory(res_ptr, res_size)?;
-        tracing::info!("Result: {:?}", result);
-        Ok(result)
-    }
-}
+pub use context::PvqExecutorContext;
+pub use error::PvqExecutorError;
+pub use executor::PvqExecutor;

pvq-extension/src/error.rs

@@ -5,6 +5,7 @@ use scale_info::prelude::fmt;
 use scale_info::prelude::fmt::{Display, Formatter};
 
 /// Errors that can occur when working with extensions
+// Typically will be used as a UserError
 #[derive(Debug)]
 pub enum ExtensionError {
     /// Permission denied for the requested operation

pvq-extension/src/executor.rs

@@ -37,9 +37,8 @@ impl<C: CallDataTuple, P: PermissionController> ExtensionsExecutor<C, P> {
     /// # Returns
     ///
     /// The result of the execution or an error
-    pub fn execute_method(&mut self, program: &[u8], args: &[u8], gas_limit: u64) -> PvqResult {
-        self.executor
-            .execute(program, args, gas_limit)
-            .map_err(|e| PvqError::Custom(scale_info::prelude::format!("{:?}", e)))
+    pub fn execute_method(&mut self, program: &[u8], args: &[u8], gas_limit: Option<i64>) -> (PvqResult, Option<i64>) {
+        let (result, gas_remaining) = self.executor.execute(program, args, gas_limit);
+        (result.map_err(PvqError::from), gas_remaining)
     }
 }

pvq-primitives/src/lib.rs

@@ -2,7 +2,7 @@
 
 extern crate alloc;
 
-use alloc::{string::String, vec::Vec};
+use alloc::vec::Vec;
 use parity_scale_codec::{Decode, Encode};
 use scale_info::TypeInfo;
 
@@ -13,7 +13,10 @@ pub enum PvqError {
     FailedToDecode,
     InvalidPvqProgramFormat,
     QueryExceedsWeightLimit,
-    Custom(String),
+    Trap,
+    MemoryAccessError,
+    HostCallError,
+    Other,
 }
 
 pub type PvqResult = Result<PvqResponse, PvqError>;

pvq-test-runner/src/lib.rs

@@ -76,7 +76,8 @@ impl TestRunner {
         program_blob: &[u8],
         input_data: &[u8],
     ) -> Result<Vec<u8>, pvq_primitives::PvqError> {
-        self.executor.execute_method(program_blob, input_data, 0)
+        let (result, _) = self.executor.execute_method(program_blob, input_data, None);
+        result
     }
 }