fix: resume program load cycles

nervosnetwork · Apr 11, 2022 · 612c58d · 612c58d
1 parent 7eaff14
commit 612c58d
Show file tree

Hide file tree

Showing 5 changed files with 190 additions and 180 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/script/Cargo.toml b/script/Cargo.toml
@@ -39,3 +39,4 @@ tiny-keccak = { version = "2.0", features = ["sha3"] }
 ckb-crypto = { path = "../util/crypto", version = "= 0.102.0-pre" }
 ckb-db-schema = { path = "../db-schema", version = "= 0.102.0-pre" }
 tempfile = "3.0"
+rand = "0.8.4"
diff --git a/script/src/types.rs b/script/src/types.rs
@@ -8,7 +8,8 @@ use ckb_vm::{
     machine::{VERSION0, VERSION1},
     memory::{FLAG_EXECUTABLE, FLAG_FREEZED},
     snapshot::{make_snapshot, Snapshot},
-    CoreMachine as _, Memory, SupportMachine, ISA_B, ISA_IMC, ISA_MOP, RISCV_PAGESIZE,
+    CoreMachine as _, Error as VMInternalError, Memory, SupportMachine, ISA_B, ISA_IMC, ISA_MOP,
+    RISCV_PAGESIZE,
 };
 use serde::{Deserialize, Serialize};
 use std::fmt;
@@ -100,33 +101,40 @@ pub(crate) type Machine<'a> = AsmMachine<'a>;
 pub(crate) type Machine<'a> = TraceMachine<'a, CoreMachine>;
 
 pub struct ResumableMachine<'a> {
-    pub(crate) machine: Machine<'a>,
-    pub(crate) program_loaded: bool,
+    machine: Machine<'a>,
+    pub(crate) program_bytes_cycles: Option<Cycle>,
 }
 
 impl<'a> ResumableMachine<'a> {
-    pub(crate) fn new(machine: Machine<'a>, program_loaded: bool) -> Self {
+    pub(crate) fn new(machine: Machine<'a>, program_bytes_cycles: Option<Cycle>) -> Self {
         ResumableMachine {
             machine,
-            program_loaded,
+            program_bytes_cycles,
         }
     }
 
     pub(crate) fn cycles(&self) -> Cycle {
         self.machine.machine.cycles()
     }
 
-    #[cfg(test)]
-    pub(crate) fn set_cycles(&mut self, cycles: Cycle) {
-        self.machine.machine.set_cycles(cycles)
-    }
-
     pub(crate) fn set_max_cycles(&mut self, cycles: Cycle) {
         set_vm_max_cycles(&mut self.machine, cycles)
     }
 
     pub fn program_loaded(&self) -> bool {
-        self.program_loaded
+        self.program_bytes_cycles.is_none()
+    }
+
+    pub fn add_cycles(&mut self, cycles: Cycle) -> Result<(), VMInternalError> {
+        self.machine.machine.add_cycles(cycles)
+    }
+
+    pub fn run(&mut self) -> Result<i8, VMInternalError> {
+        if let Some(cycles) = self.program_bytes_cycles {
+            self.add_cycles(cycles)?;
+            self.program_bytes_cycles = None;
+        }
+        self.machine.run()
     }
 }
 
@@ -275,7 +283,7 @@ impl TryFrom<TransactionState<'_>> for TransactionSnapshot {
 
         let (snap, current_cycles) = if let Some(mut vm) = vm {
             // we should not capture snapshot if load program failed by exceeded cycles
-            if vm.program_loaded {
+            if vm.program_loaded() {
                 let vm_cycles = vm.cycles();
                 // To be consistent with the mainnet, add this flag to enable this behavior after hardfork
                 if !enable_backup_page_flags {

diff --git a/script/src/verify.rs b/script/src/verify.rs
@@ -669,7 +669,7 @@ impl<'a, DL: CellDataProvider + HeaderProvider> TransactionScriptsVerifier<'a, D
 
         if let Some(mut vm) = vm {
             vm.set_max_cycles(limit_cycles);
-            match vm.machine.run() {
+            match vm.run() {
                 Ok(code) => {
                     self.tracing_data_as_code_pages.borrow_mut().clear();
                     if code == 0 {
@@ -1017,9 +1017,13 @@ impl<'a, DL: CellDataProvider + HeaderProvider> TransactionScriptsVerifier<'a, D
             let bytes = machine
                 .load_program(&program, &[])
                 .map_err(map_vm_internal_error)?;
-            let load_ret = machine.machine.add_cycles(transferred_byte_cycles(bytes));
+            let program_bytes_cycles = transferred_byte_cycles(bytes);
+            let load_ret = machine.machine.add_cycles(program_bytes_cycles);
             if matches!(load_ret, Err(ref error) if error == &VMInternalError::CyclesExceeded) {
-                return Ok(ChunkState::suspended(ResumableMachine::new(machine, false)));
+                return Ok(ChunkState::suspended(ResumableMachine::new(
+                    machine,
+                    Some(program_bytes_cycles),
+                )));
             }
             load_ret.map_err(|e| ScriptError::VMInternalError(format!("{:?}", e)))?;
         }
@@ -1035,7 +1039,7 @@ impl<'a, DL: CellDataProvider + HeaderProvider> TransactionScriptsVerifier<'a, D
             }
             Err(error) => match error {
                 VMInternalError::CyclesExceeded => {
-                    Ok(ChunkState::suspended(ResumableMachine::new(machine, true)))
+                    Ok(ChunkState::suspended(ResumableMachine::new(machine, None)))
                 }
                 _ => {
                     self.tracing_data_as_code_pages.borrow_mut().clear();