Switch to a unified pool instead of per-thread one

crates/services/executor/src/executor.rs

@@ -338,6 +338,7 @@ pub struct BlockExecutor<R> {
     relayer: R,
     consensus_params: ConsensusParameters,
     options: ExecutionOptions,
+    vm_pool: vm_pool::VmPool,
 }
 
 impl<R> BlockExecutor<R> {
@@ -350,6 +351,7 @@ impl<R> BlockExecutor<R> {
             relayer,
             consensus_params,
             options,
+            vm_pool: Default::default(),
         })
     }
 }
@@ -1433,15 +1435,15 @@ where
         );
 
         // We get memory from the pool and are careful to return it there when done
-        *vm.memory_mut() = vm_pool::get_vm_memory();
+        *vm.memory_mut() = self.vm_pool.get_new();
 
         // This returns the VM memory to the pool to before returning
         macro_rules! recycling_try {
             ($expr:expr) => {
                 match $expr {
                     Ok(val) => val,
                     Err(e) => {
-                        vm_pool::recycle_vm_memory(mem::take(vm.memory_mut()));
+                        self.vm_pool.recycle(mem::take(vm.memory_mut()));
                         return Err(e.into())
                     }
                 }
@@ -1477,11 +1479,11 @@ where
         // only commit state changes if execution was a success
         if !reverted {
             self.log_backtrace(&vm, &receipts);
-            vm_pool::recycle_vm_memory(mem::take(vm.memory_mut()));
+            self.vm_pool.recycle(mem::take(vm.memory_mut()));
             let changes = sub_block_db_commit.into_changes();
             storage_tx.commit_changes(changes)?;
         } else {
-            vm_pool::recycle_vm_memory(mem::take(vm.memory_mut()));
+            self.vm_pool.recycle(mem::take(vm.memory_mut()));
         }
 
         self.update_tx_outputs(storage_tx, tx_id, &mut tx)?;

crates/services/executor/src/vm_pool.rs

@@ -1,39 +1,52 @@
-use std::cell::RefCell;
+use core::fmt;
+use std::sync::{
+    Arc,
+    Mutex,
+};
 
 use fuel_core_types::fuel_vm::interpreter::Memory;
 
-thread_local! {
-    static POOL: RefCell<Vec<Memory>> = RefCell::new(Vec::new());
+#[derive(Default, Clone)]
+pub(crate) struct VmPool {
+    pool: Arc<Mutex<Vec<Memory>>>,
 }
-
-/// Gets a new VM memory instance from the pool.
-pub(crate) fn get_vm_memory() -> Memory {
-    POOL.with(|pool| {
-        let mut pool = pool.borrow_mut();
+impl VmPool {
+    /// Gets a new VM memory instance from the pool.
+    pub(crate) fn get_new(&self) -> Memory {
+        let mut pool = self.pool.lock().expect("poisoned");
         pool.pop().unwrap_or_else(Memory::new)
-    })
-}
+    }
 
-/// Recycles a VM memory instance back into the pool.
-pub(crate) fn recycle_vm_memory(mut mem: Memory) {
-    POOL.with(|pool| {
+    /// Recycles a VM memory instance back into the pool.
+    pub(crate) fn recycle(&self, mut mem: Memory) {
         mem.reset();
-        let mut pool = pool.borrow_mut();
+        let mut pool = self.pool.lock().expect("poisoned");
         pool.push(mem);
-    })
+    }
+}
+
+impl fmt::Debug for VmPool {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self.pool.lock() {
+            Ok(pool) => write!(f, "VmPool {{ pool: [{} items] }}", pool.len()),
+            Err(_) => write!(f, "VmPool {{ pool: [poisoned] }}"),
+        }
+    }
 }
 
 #[test]
 fn test_vm_pool() {
-    let mut mem = get_vm_memory();
+    let pool = VmPool::default();
+
+    let mut mem = pool.get_new();
     mem.grow_stack(1024).expect("Unable to grow stack");
     mem.write_bytes_noownerchecks(0, [1, 2, 3, 4])
         .expect("Unable to write stack");
     let ptr1 = mem.stack_raw() as *const _ as *const u8 as usize;
-    recycle_vm_memory(mem);
+    pool.recycle(mem);
 
     // Make sure we get the same memory allocation back
-    let mem = get_vm_memory();
+    let mem = pool.get_new();
     let ptr2 = mem.stack_raw() as *const _ as *const u8 as usize;
     assert_eq!(ptr1, ptr2);
 }