From 4778ddda168de96934e5e18bd06644d0b2cfe00c Mon Sep 17 00:00:00 2001
From: Roman Maslennikov <m3angreen@gmail.com>
Date: Tue, 27 Aug 2024 02:29:52 +0400
Subject: [PATCH] fix(benchmarks,runtime-interface): clear store by counter
 when running benchmarks

---
 pallets/gear/Cargo.toml                 |  1 +
 runtime-interface/Cargo.toml            |  3 +++
 runtime-interface/sandbox/Cargo.toml    |  1 +
 runtime-interface/sandbox/src/detail.rs | 28 +++++++++++++++++++++++++
 4 files changed, 33 insertions(+)

diff --git a/pallets/gear/Cargo.toml b/pallets/gear/Cargo.toml
index 99284159692..af9fb543119 100644
--- a/pallets/gear/Cargo.toml
+++ b/pallets/gear/Cargo.toml
@@ -182,6 +182,7 @@ runtime-benchmarks = [
     "frame-system/runtime-benchmarks",
     "frame-support/runtime-benchmarks",
     "common/runtime-benchmarks",
+    "gear-runtime-interface/runtime-benchmarks",
     "gear-core-backend/mock",
     "gear-core-errors/codec",
     "gear-sandbox",
diff --git a/runtime-interface/Cargo.toml b/runtime-interface/Cargo.toml
index f58660e1ac9..01c4ad8e4b6 100644
--- a/runtime-interface/Cargo.toml
+++ b/runtime-interface/Cargo.toml
@@ -53,3 +53,6 @@ std = [
 	"ark-scale/std",
 	"sha2/std"
 ]
+runtime-benchmarks = [
+	"gear-sandbox-interface/runtime-benchmarks"
+]
\ No newline at end of file
diff --git a/runtime-interface/sandbox/Cargo.toml b/runtime-interface/sandbox/Cargo.toml
index f34553be354..8e74c327f55 100644
--- a/runtime-interface/sandbox/Cargo.toml
+++ b/runtime-interface/sandbox/Cargo.toml
@@ -28,3 +28,4 @@ std = [
   "sp-runtime-interface/std",
   "sp-wasm-interface/std",
 ]
+runtime-benchmarks = []
\ No newline at end of file
diff --git a/runtime-interface/sandbox/src/detail.rs b/runtime-interface/sandbox/src/detail.rs
index 6f67076c211..c4f401d7f6b 100644
--- a/runtime-interface/sandbox/src/detail.rs
+++ b/runtime-interface/sandbox/src/detail.rs
@@ -47,6 +47,18 @@ impl Sandboxes {
 
         &mut self.store
     }
+
+    // Clears the underlying store if the counter exceeds the limit.
+    #[cfg(feature = "runtime-benchmarks")]
+    pub fn clear(&mut self) {
+        BENCH_SANDBOX_RESET_COUNTER.with_borrow_mut(|c| {
+            if *c >= BENCH_SANDBOX_RESET_COUNTER_LIMIT {
+                *c = 0;
+                self.store.clear();
+            }
+            *c += 1;
+        });
+    }
 }
 
 thread_local! {
@@ -453,6 +465,14 @@ pub fn memory_new(context: &mut dyn FunctionContext, initial: u32, maximum: u32)
 
         let data_ptr: *const _ = caller.data();
         method_result = SANDBOXES.with(|sandboxes| {
+            // The usual method to clear the store doesn't work for benchmarks (see `Sanboxes::get`).
+            // This issue is more prominent in so-called "onetime syscall" benchmarks because
+            // they run with minimal steps and a large number of repeats, leading to significant slowdowns.
+            // Therefore, we have to clear it manually if the `BENCH_SANDBOX_RESET_COUNTER` exceeds the limit.
+            // Otherwise, the store becomes too big and will slow down the benchmarks.
+            #[cfg(feature = "runtime-benchmarks")]
+            sandboxes.borrow_mut().clear();
+
             sandboxes
                 .borrow_mut()
                 .get(data_ptr as usize)
@@ -577,3 +597,11 @@ pub fn set_global_val(
 
     method_result
 }
+
+#[cfg(feature = "runtime-benchmarks")]
+const BENCH_SANDBOX_RESET_COUNTER_LIMIT: u32 = 100;
+
+#[cfg(feature = "runtime-benchmarks")]
+thread_local! {
+    static BENCH_SANDBOX_RESET_COUNTER: RefCell<u32> = const { RefCell::new(0) };
+}