Fix task distribution (#1013)

Co-authored-by: Metadata Update Bot <[email protected]>

pallets/tasks/src/lib.rs

@@ -658,7 +658,7 @@ pub mod pallet {
 
 		/// Sets the reward for write tasks.
 		///
-		// # Flow
+		/// # Flow
 		///    1. Ensure the origin of the transaction is a root user.
 		///    2. Insert the new reward amount for the specified network into the [`NetworkWriteReward`] storage.
 		///    3. Emit an event [`Event::WriteTaskRewardSet`] indicating the write task reward has been set.
@@ -1275,7 +1275,7 @@ pub mod pallet {
 		///     - Retrieve all shards associated with the `network`.
 		///     - Collect these shards into a list.
 		///     - Sort the list based on the number of tasks each shard currently has.
-		///     - Iterate through the sorted list in reverse order.
+		///     - Iterate through the sorted list.
 		///     - For each shard in the list, call [`Self::schedule_tasks_shard`] to schedule tasks.
 		fn schedule_tasks(network: NetworkId, shard_id: Option<ShardId>) {
 			if let Some(shard_id) = shard_id {
@@ -1287,7 +1287,7 @@ pub mod pallet {
 						.count()
 						.cmp(&ShardTasks::<T>::iter_prefix(*b).count())
 				});
-				shards.into_iter().rev().for_each(|(shard, _)| {
+				shards.into_iter().for_each(|(shard, _)| {
 					Self::schedule_tasks_shard(network, shard);
 				});
 			}

pallets/tasks/src/tests.rs

@@ -1813,6 +1813,98 @@ fn test_task_execution_order() {
 	});
 }
 
+#[test]
+fn test_multi_shard_distribution() {
+	new_test_ext().execute_with(|| {
+		// Shard creation
+		for i in 0..3 {
+			Shards::create_shard(
+				ETHEREUM,
+				[[0u8; 32].into(), [1u8; 32].into(), [2u8; 32].into()].to_vec(),
+				1,
+			);
+			ShardState::<Test>::insert(i, ShardStatus::Online);
+			Tasks::shard_online(i, ETHEREUM);
+		}
+
+		// Tasks creation and assingment
+		for _ in 0..9 {
+			assert_ok!(Tasks::create_task(
+				RawOrigin::Signed([0; 32].into()).into(),
+				mock_task(ETHEREUM, 3)
+			));
+		}
+
+		assert_eq!(ShardTasks::<Test>::iter_prefix(0).count(), 3);
+		assert_eq!(ShardTasks::<Test>::iter_prefix(1).count(), 3);
+		assert_eq!(ShardTasks::<Test>::iter_prefix(2).count(), 3);
+	});
+}
+
+#[test]
+fn test_multi_shard_distribution_task_more_than_limit() {
+	new_test_ext().execute_with(|| {
+		// Shard creation
+		for i in 0..3 {
+			Shards::create_shard(
+				ETHEREUM,
+				[[0u8; 32].into(), [1u8; 32].into(), [2u8; 32].into()].to_vec(),
+				1,
+			);
+			ShardState::<Test>::insert(i, ShardStatus::Online);
+			Tasks::shard_online(i, ETHEREUM);
+		}
+
+		assert_ok!(Tasks::set_shard_task_limit(RawOrigin::Root.into(), ETHEREUM, 5));
+
+		// Tasks creation and assingment
+		for _ in 0..30 {
+			assert_ok!(Tasks::create_task(
+				RawOrigin::Signed([0; 32].into()).into(),
+				mock_task(ETHEREUM, 3)
+			));
+		}
+
+		assert_eq!(ShardTasks::<Test>::iter_prefix(0).count(), 5);
+		assert_eq!(ShardTasks::<Test>::iter_prefix(1).count(), 5);
+		assert_eq!(ShardTasks::<Test>::iter_prefix(2).count(), 5);
+	});
+}
+
+#[test]
+fn test_multi_shard_distribution_task_before_shard_online() {
+	new_test_ext().execute_with(|| {
+		// Shard creation
+		for i in 0..3 {
+			Shards::create_shard(
+				ETHEREUM,
+				[[0u8; 32].into(), [1u8; 32].into(), [2u8; 32].into()].to_vec(),
+				1,
+			);
+			ShardState::<Test>::insert(i, ShardStatus::Online);
+		}
+
+		assert_ok!(Tasks::set_shard_task_limit(RawOrigin::Root.into(), ETHEREUM, 10));
+
+		// Tasks creation and assingment
+		for _ in 0..25 {
+			assert_ok!(Tasks::create_task(
+				RawOrigin::Signed([0; 32].into()).into(),
+				mock_task(ETHEREUM, 3)
+			));
+		}
+
+		// shards come online when there are already some pending tasks to work with
+		for i in 0..3 {
+			Tasks::shard_online(i, ETHEREUM);
+		}
+
+		assert_eq!(ShardTasks::<Test>::iter_prefix(0).count(), 10);
+		assert_eq!(ShardTasks::<Test>::iter_prefix(1).count(), 10);
+		assert_eq!(ShardTasks::<Test>::iter_prefix(2).count(), 5);
+	});
+}
+
 #[test]
 fn test_assingment_with_diff_shard_size() {
 	new_test_ext().execute_with(|| {