feat(common, gtest): add AuxiliaryTaskpool (#4141)

common/src/auxiliary/mailbox.rs

@@ -17,10 +17,8 @@
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
 //! Auxiliary implementation of the mailbox.
-
-use super::AuxiliaryDoubleStorageWrap;
 use crate::{
-    auxiliary::DoubleBTreeMap,
+    auxiliary::{AuxiliaryDoubleStorageWrap, BlockNumber, DoubleBTreeMap},
     storage::{Interval, MailboxError, MailboxImpl, MailboxKeyGen},
 };
 use core::cell::RefCell;
@@ -39,8 +37,7 @@ pub type AuxiliaryMailbox<MailboxCallbacks> = MailboxImpl<
     MailboxCallbacks,
     MailboxKeyGen<ProgramId>,
 >;
-/// Block number type.
-pub type BlockNumber = u32;
+
 /// Type represents message stored in the mailbox.
 pub type MailboxedMessage = UserStoredMessage;
 

common/src/auxiliary/mod.rs

@@ -22,13 +22,17 @@
 
 pub mod gas_provider;
 pub mod mailbox;
+pub mod taskpool;
 
 use crate::storage::{
     Counted, CountedByKey, DoubleMapStorage, GetFirstPos, GetSecondPos, IterableByKeyMap,
     IteratorWrap, KeyIterableByKeyMap, MapStorage,
 };
 use alloc::collections::btree_map::{BTreeMap, Entry, IntoIter};
 
+/// An "auxiliary" block number type.
+pub type BlockNumber = u32;
+
 /// Double key `BTreeMap`.
 ///
 /// Basically is just a map of the map.

common/src/auxiliary/taskpool.rs

@@ -0,0 +1,58 @@
+//! Auxiliary implementation of the taskpool.
+
+use super::{AuxiliaryDoubleStorageWrap, BlockNumber, DoubleBTreeMap};
+use crate::scheduler::{ScheduledTask, TaskPoolImpl};
+use gear_core::ids::ProgramId;
+use std::cell::RefCell;
+
+pub type AuxiliaryTaskpool<TaskPoolCallbacks> = TaskPoolImpl<
+    TaskPoolStorageWrap,
+    ScheduledTask<ProgramId>,
+    TaskPoolErrorImpl,
+    TaskPoolErrorImpl,
+    TaskPoolCallbacks,
+>;
+
+std::thread_local! {
+    pub(crate) static TASKPOOL_STORAGE: RefCell<DoubleBTreeMap<BlockNumber, ScheduledTask<ProgramId>, ()>> = const { RefCell::new(DoubleBTreeMap::new()) };
+}
+
+pub struct TaskPoolStorageWrap;
+
+impl AuxiliaryDoubleStorageWrap for TaskPoolStorageWrap {
+    type Key1 = BlockNumber;
+    type Key2 = ScheduledTask<ProgramId>;
+    type Value = ();
+
+    fn with_storage<F, R>(f: F) -> R
+    where
+        F: FnOnce(&DoubleBTreeMap<Self::Key1, Self::Key2, Self::Value>) -> R,
+    {
+        TASKPOOL_STORAGE.with_borrow(f)
+    }
+
+    fn with_storage_mut<F, R>(f: F) -> R
+    where
+        F: FnOnce(&mut DoubleBTreeMap<Self::Key1, Self::Key2, Self::Value>) -> R,
+    {
+        TASKPOOL_STORAGE.with_borrow_mut(f)
+    }
+}
+
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub enum TaskPoolErrorImpl {
+    /// Occurs when given task already exists in task pool.
+    DuplicateTask,
+    /// Occurs when task wasn't found in storage.
+    TaskNotFound,
+}
+
+impl crate::scheduler::TaskPoolError for TaskPoolErrorImpl {
+    fn duplicate_task() -> Self {
+        Self::DuplicateTask
+    }
+
+    fn task_not_found() -> Self {
+        Self::TaskNotFound
+    }
+}

gtest/src/lib.rs

@@ -452,6 +452,7 @@ mod mailbox;
 mod manager;
 mod program;
 mod system;
+mod taskpool;
 
 pub use crate::log::{BlockRunResult, CoreLog, Log};
 pub use codec;

gtest/src/mailbox/actor.rs

@@ -18,7 +18,10 @@
 
 use crate::{manager::ExtManager, Log, GAS_ALLOWANCE};
 use codec::Encode;
-use gear_common::{auxiliary::mailbox::*, storage::Interval};
+use gear_common::{
+    auxiliary::{mailbox::*, BlockNumber},
+    storage::Interval,
+};
 use gear_core::{
     ids::{prelude::MessageIdExt, MessageId, ProgramId},
     message::{ReplyMessage, ReplyPacket},

gtest/src/mailbox/manager.rs

@@ -20,7 +20,7 @@
 
 use crate::blocks::BlocksManager;
 use gear_common::{
-    auxiliary::mailbox::*,
+    auxiliary::{mailbox::*, BlockNumber},
     storage::{GetCallback, Interval, IterableByKeyMap, Mailbox, MailboxCallbacks},
 };
 use gear_core::ids::{MessageId, ProgramId};

gtest/src/taskpool.rs

@@ -0,0 +1,133 @@
+use gear_common::{
+    auxiliary::{
+        taskpool::{AuxiliaryTaskpool, TaskPoolErrorImpl, TaskPoolStorageWrap},
+        BlockNumber,
+    },
+    scheduler::{ScheduledTask, TaskPool, TaskPoolCallbacks},
+    storage::KeyIterableByKeyMap,
+    ProgramId,
+};
+
+/// Task pool manager for gtest environment.
+///
+/// TODO(ap): wait for #4119 and work on integrating this into it, until then
+/// allow(dead_code).
+#[derive(Debug, Default)]
+#[allow(dead_code)]
+pub(crate) struct TaskPoolManager;
+
+#[allow(dead_code)]
+impl TaskPoolManager {
+    pub(crate) fn add(
+        &self,
+        block_number: BlockNumber,
+        task: ScheduledTask<ProgramId>,
+    ) -> Result<(), TaskPoolErrorImpl> {
+        <AuxiliaryTaskpool<TaskPoolCallbacksImpl> as TaskPool>::add(block_number, task)
+    }
+
+    pub(crate) fn clear(&self) {
+        <AuxiliaryTaskpool<TaskPoolCallbacksImpl> as TaskPool>::clear();
+    }
+
+    pub(crate) fn contains(
+        &self,
+        block_number: &BlockNumber,
+        task: &ScheduledTask<ProgramId>,
+    ) -> bool {
+        <AuxiliaryTaskpool<TaskPoolCallbacksImpl> as TaskPool>::contains(block_number, task)
+    }
+
+    pub(crate) fn delete(
+        &self,
+        block_number: BlockNumber,
+        task: ScheduledTask<ProgramId>,
+    ) -> Result<(), TaskPoolErrorImpl> {
+        <AuxiliaryTaskpool<TaskPoolCallbacksImpl> as TaskPool>::delete(block_number, task)
+    }
+
+    pub(crate) fn drain(
+        &self,
+        block_number: BlockNumber,
+    ) -> <TaskPoolStorageWrap as KeyIterableByKeyMap>::DrainIter {
+        AuxiliaryTaskpool::<TaskPoolCallbacksImpl>::drain_prefix_keys(block_number)
+    }
+}
+
+pub(crate) struct TaskPoolCallbacksImpl;
+
+impl TaskPoolCallbacks for TaskPoolCallbacksImpl {
+    type OnAdd = ();
+    type OnDelete = ();
+}
+
+#[cfg(test)]
+mod tests {
+    use gear_common::{scheduler::ScheduledTask, ProgramId};
+
+    use super::TaskPoolManager;
+
+    #[test]
+    fn test_taskpool() {
+        let manager = TaskPoolManager;
+
+        let block_1_tasks = [
+            ScheduledTask::<ProgramId>::SendDispatch(42.into()),
+            ScheduledTask::<ProgramId>::SendUserMessage {
+                message_id: 422.into(),
+                to_mailbox: true,
+            },
+        ];
+        let block_2_tasks = [
+            ScheduledTask::<ProgramId>::RemoveGasReservation(922.into(), 1.into()),
+            ScheduledTask::<ProgramId>::RemoveFromWaitlist(42.into(), 44.into()),
+        ];
+
+        block_1_tasks
+            .iter()
+            .for_each(|task| manager.add(1, task.clone()).unwrap());
+
+        block_2_tasks
+            .iter()
+            .for_each(|task| manager.add(2, task.clone()).unwrap());
+
+        for task in block_1_tasks.iter() {
+            assert!(manager.contains(&1, task));
+        }
+
+        for task in block_2_tasks.iter() {
+            assert!(manager.contains(&2, task));
+        }
+
+        for task in manager.drain(1) {
+            assert!(
+                block_1_tasks.contains(&task),
+                "task not found in block 1 tasks"
+            );
+        }
+
+        for task in manager.drain(2) {
+            assert!(
+                block_2_tasks.contains(&task),
+                "task not found in block 2 tasks"
+            );
+        }
+
+        for task in block_1_tasks.iter() {
+            assert!(!manager.contains(&1, task));
+        }
+
+        for task in block_2_tasks.iter() {
+            assert!(!manager.contains(&2, task));
+        }
+
+        let task = ScheduledTask::<ProgramId>::RemoveFromMailbox(422.into(), 16.into());
+        manager.add(3, task.clone()).unwrap();
+        manager.add(4, task.clone()).unwrap();
+        manager.delete(4, task.clone()).unwrap();
+        manager.clear();
+
+        assert!(!manager.contains(&3, &task));
+        assert!(!manager.contains(&4, &task));
+    }
+}