feat(ethexe): sparse memory pages (#4390)

Makefile

@@ -6,7 +6,7 @@ ethexe-pre-commit: ethexe-contracts-pre-commit ethexe-pre-commit-no-contracts
 ethexe-pre-commit-no-contracts:
 	@ echo " > Formatting ethexe" && cargo +nightly fmt --all -- --config imports_granularity=Crate,edition=2021
 	@ echo " >> Clippy checking ethexe" && cargo clippy -p "ethexe-*" --all-targets --all-features -- --no-deps -D warnings
-	@ echo " >>> Testing ethexe" && cargo test -p "ethexe-*"
+	@ echo " >>> Testing ethexe" && cargo nextest run -p "ethexe-*"
 
 # Building ethexe contracts
 .PHONY: ethexe-contracts-pre-commit

ethexe/db/src/database.rs

@@ -28,8 +28,8 @@ use ethexe_common::{
     gear::StateTransition,
 };
 use ethexe_runtime_common::state::{
-    Allocations, DispatchStash, HashOf, Mailbox, MemoryPages, MessageQueue, ProgramState, Storage,
-    Waitlist,
+    Allocations, DispatchStash, HashOf, Mailbox, MemoryPages, MemoryPagesRegion, MessageQueue,
+    ProgramState, Storage, Waitlist,
 };
 use gear_core::{
     code::InstrumentedCode,
@@ -531,10 +531,21 @@ impl Storage for Database {
         })
     }
 
+    fn read_pages_region(&self, hash: HashOf<MemoryPagesRegion>) -> Option<MemoryPagesRegion> {
+        self.cas.read(&hash.hash()).map(|data| {
+            MemoryPagesRegion::decode(&mut &data[..])
+                .expect("Failed to decode data into `MemoryPagesRegion`")
+        })
+    }
+
     fn write_pages(&self, pages: MemoryPages) -> HashOf<MemoryPages> {
         unsafe { HashOf::new(self.cas.write(&pages.encode())) }
     }
 
+    fn write_pages_region(&self, pages_region: MemoryPagesRegion) -> HashOf<MemoryPagesRegion> {
+        unsafe { HashOf::new(self.cas.write(&pages_region.encode())) }
+    }
+
     fn read_allocations(&self, hash: HashOf<Allocations>) -> Option<Allocations> {
         self.cas.read(&hash.hash()).map(|data| {
             Allocations::decode(&mut &data[..]).expect("Failed to decode data into `Allocations`")

ethexe/processor/src/host/threads.rs

@@ -22,7 +22,10 @@ use crate::Database;
 use core::fmt;
 use ethexe_db::BlockMetaStorage;
 use ethexe_runtime_common::{
-    state::{ActiveProgram, HashOf, Program, ProgramState, Storage},
+    state::{
+        ActiveProgram, HashOf, MemoryPages, MemoryPagesInner, MemoryPagesRegionInner, Program,
+        ProgramState, RegionIdx, Storage,
+    },
     BlockInfo,
 };
 use gear_core::{ids::ProgramId, memory::PageBuf, pages::GearPage};
@@ -42,7 +45,8 @@ pub struct ThreadParams {
     pub db: Database,
     pub block_info: BlockInfo,
     pub state_hash: H256,
-    pub pages: Option<BTreeMap<GearPage, HashOf<PageBuf>>>,
+    pages_registry_cache: Option<MemoryPagesInner>,
+    pages_regions_cache: Option<BTreeMap<RegionIdx, MemoryPagesRegionInner>>,
 }
 
 impl fmt::Debug for ThreadParams {
@@ -52,19 +56,38 @@ impl fmt::Debug for ThreadParams {
 }
 
 impl ThreadParams {
-    pub fn pages(&mut self) -> &BTreeMap<GearPage, HashOf<PageBuf>> {
-        self.pages.get_or_insert_with(|| {
+    pub fn get_page_region(
+        &mut self,
+        page: GearPage,
+    ) -> Option<&BTreeMap<GearPage, HashOf<PageBuf>>> {
+        let pages_registry = self.pages_registry_cache.get_or_insert_with(|| {
             let ProgramState {
                 program: Program::Active(ActiveProgram { pages_hash, .. }),
                 ..
             } = self.db.read_state(self.state_hash).expect(UNKNOWN_STATE)
             else {
-                // TODO: consider me.
-                panic!("Couldn't get pages hash for inactive program!")
+                unreachable!("program that is currently running can't be inactive");
             };
 
             pages_hash.query(&self.db).expect(UNKNOWN_STATE).into()
-        })
+        });
+
+        let region_idx = MemoryPages::page_region(page);
+
+        let region_hash = pages_registry.get(&region_idx)?;
+
+        let pages_regions = self
+            .pages_regions_cache
+            .get_or_insert_with(Default::default);
+
+        let page_region = pages_regions.entry(region_idx).or_insert_with(|| {
+            self.db
+                .read_pages_region(*region_hash)
+                .expect("Pages region not found")
+                .into()
+        });
+
+        Some(&*page_region)
     }
 }
 
@@ -92,11 +115,11 @@ pub fn set(db: Database, chain_head: H256, state_hash: H256) {
             timestamp: header.timestamp,
         },
         state_hash,
-        pages: None,
+        pages_registry_cache: None,
+        pages_regions_cache: None,
     }))
 }
 
-// TODO: consider Database mutability.
 pub fn with_db<T>(f: impl FnOnce(&Database) -> T) -> T {
     PARAMS.with_borrow(|v| {
         let params = v.as_ref().expect(UNSET_PANIC);
@@ -113,16 +136,6 @@ pub fn chain_head_info() -> BlockInfo {
     })
 }
 
-// TODO: consider Database mutability.
-#[allow(unused)]
-pub fn with_db_and_state_hash<T>(f: impl FnOnce(&Database, H256) -> T) -> T {
-    PARAMS.with_borrow(|v| {
-        let params = v.as_ref().expect(UNSET_PANIC);
-
-        f(&params.db, params.state_hash)
-    })
-}
-
 pub fn with_params<T>(f: impl FnOnce(&mut ThreadParams) -> T) -> T {
     PARAMS.with_borrow_mut(|v| {
         let params = v.as_mut().expect(UNSET_PANIC);
@@ -139,7 +152,7 @@ impl LazyPagesStorage for EthexeHostLazyPages {
         with_params(|params| {
             let page = PageKey::page_from_buf(key);
 
-            let page_hash = params.pages().get(&page).cloned()?;
+            let page_hash = params.get_page_region(page)?.get(&page).cloned()?;
 
             let data = params.db.read_page_data(page_hash).expect("Page not found");
 
@@ -153,7 +166,10 @@ impl LazyPagesStorage for EthexeHostLazyPages {
         with_params(|params| {
             let page = PageKey::page_from_buf(key);
 
-            params.pages().contains_key(&page)
+            params
+                .get_page_region(page)
+                .map(|region| region.contains_key(&page))
+                .unwrap_or(false)
         })
     }
 }

ethexe/runtime/common/src/journal.rs

@@ -323,7 +323,7 @@ impl<S: Storage> JournalHandler for Handler<'_, S> {
             };
 
             pages_hash.modify_pages(storage, |pages| {
-                pages.update(storage.write_pages_data(pages_data));
+                pages.update_and_store_regions(storage, storage.write_pages_data(pages_data));
             });
 
             Ok(())
@@ -351,7 +351,7 @@ impl<S: Storage> JournalHandler for Handler<'_, S> {
 
                 if !removed_pages.is_empty() {
                     pages_hash.modify_pages(storage, |pages| {
-                        pages.remove(&removed_pages);
+                        pages.remove_and_store_regions(storage, &removed_pages);
                     })
                 }
             });

ethexe/runtime/common/src/lib.rs

@@ -22,7 +22,7 @@
 
 extern crate alloc;
 
-use alloc::{collections::BTreeMap, vec::Vec};
+use alloc::vec::Vec;
 use core_processor::{
     common::{ExecutableActorData, JournalNote},
     configs::{BlockConfig, SyscallName},
@@ -31,14 +31,12 @@ use core_processor::{
 use gear_core::{
     code::{InstrumentedCode, MAX_WASM_PAGES_AMOUNT},
     ids::ProgramId,
-    memory::PageBuf,
     message::{DispatchKind, IncomingDispatch, IncomingMessage},
-    pages::GearPage,
 };
 use gear_lazy_pages_common::LazyPagesInterface;
 use gprimitives::CodeId;
 use gsys::{GasMultiplier, Percent};
-use state::{Dispatch, HashOf, ProgramState, Storage};
+use state::{Dispatch, ProgramState, Storage};
 
 pub use core_processor::configs::BlockInfo;
 pub use journal::Handler as JournalHandler;
@@ -59,7 +57,7 @@ pub trait RuntimeInterface<S: Storage> {
     type LazyPages: LazyPagesInterface + 'static;
 
     fn block_info(&self) -> BlockInfo;
-    fn init_lazy_pages(&self, pages_map: BTreeMap<GearPage, HashOf<PageBuf>>);
+    fn init_lazy_pages(&self);
     fn random_data(&self) -> (Vec<u8>, u32);
     fn storage(&self) -> &S;
 }
@@ -236,11 +234,6 @@ where
         Err(journal) => return journal,
     };
 
-    let pages_map = active_state.pages_hash.with_hash_or_default(|hash| {
-        ri.storage()
-            .read_pages(hash)
-            .expect("Cannot get memory pages")
-    });
     let actor_data = ExecutableActorData {
         allocations: allocations.into(),
         code_id,
@@ -271,7 +264,7 @@ where
 
     let random_data = ri.random_data();
 
-    ri.init_lazy_pages(pages_map.into());
+    ri.init_lazy_pages();
 
     core_processor::process::<Ext<RI::LazyPages>>(&block_config, execution_context, random_data)
         .unwrap_or_else(|err| unreachable!("{err}"))