From 68426cc10a6239a42170778d38bb2e76095aa5f1 Mon Sep 17 00:00:00 2001
From: zjb0807 <zjb0807@qq.com>
Date: Wed, 11 Jan 2023 10:21:13 +1300
Subject: [PATCH] Fix payments on-idle (#868)

* fix payments on-idle

* update comment
---
 payments/src/lib.rs   | 15 +++++++------
 payments/src/tests.rs | 52 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 59 insertions(+), 8 deletions(-)

diff --git a/payments/src/lib.rs b/payments/src/lib.rs
index 8a66cbee6..9f14d1927 100644
--- a/payments/src/lib.rs
+++ b/payments/src/lib.rs
@@ -214,16 +214,17 @@ pub mod pallet {
 		/// Hook that execute when there is leftover space in a block
 		/// This function will look for any pending scheduled tasks that can
 		/// be executed and will process them.
-		fn on_idle(now: T::BlockNumber, mut remaining_weight: Weight) -> Weight {
+		fn on_idle(now: T::BlockNumber, remaining_weight: Weight) -> Weight {
 			const MAX_TASKS_TO_PROCESS: usize = 5;
-			// reduce the weight used to read the task list
-			remaining_weight = remaining_weight.saturating_sub(T::WeightInfo::remove_task());
+			// used to read the task list
+			let mut used_weight = T::WeightInfo::remove_task();
 			let cancel_weight = T::WeightInfo::cancel();
 
 			// calculate count of tasks that can be processed with remaining weight
 			let possible_task_count: usize = remaining_weight
-				.ref_time()
+				.saturating_sub(used_weight)
 				.saturating_div(cancel_weight.ref_time())
+				.ref_time()
 				.try_into()
 				.unwrap_or(MAX_TASKS_TO_PROCESS);
 
@@ -239,9 +240,9 @@ pub mod pallet {
 				// order by oldest task to process
 				task_list.sort_by(|(_, t), (_, x)| x.when.cmp(&t.when));
 
-				while !task_list.is_empty() && remaining_weight.all_gte(cancel_weight) {
+				while !task_list.is_empty() && used_weight.all_lte(remaining_weight) {
 					if let Some((account_pair, _)) = task_list.pop() {
-						remaining_weight = remaining_weight.saturating_sub(cancel_weight);
+						used_weight = used_weight.saturating_add(cancel_weight);
 						// remove the task form the tasks storage
 						tasks.remove(&account_pair);
 
@@ -268,7 +269,7 @@ pub mod pallet {
 					}
 				}
 			});
-			remaining_weight
+			used_weight
 		}
 	}
 
diff --git a/payments/src/tests.rs b/payments/src/tests.rs
index ddcc3fa79..ddd2bf7ab 100644
--- a/payments/src/tests.rs
+++ b/payments/src/tests.rs
@@ -1,9 +1,10 @@
 use crate::{
 	mock::*,
 	types::{PaymentDetail, PaymentState},
+	weights::WeightInfo,
 	Payment as PaymentStore, PaymentHandler, ScheduledTask, ScheduledTasks, Task,
 };
-use frame_support::{assert_noop, assert_ok, storage::with_transaction};
+use frame_support::{assert_noop, assert_ok, storage::with_transaction, traits::OnIdle, weights::Weight};
 use orml_traits::MultiCurrency;
 use sp_runtime::{Percent, TransactionOutcome};
 
@@ -1398,3 +1399,52 @@ fn test_automatic_refund_works_for_multiple_payments() {
 		assert_eq!(Tokens::free_balance(CURRENCY_ID, &PAYMENT_RECIPENT_TWO), 0);
 	});
 }
+
+#[test]
+fn on_idle_works() {
+	new_test_ext().execute_with(|| {
+		assert_eq!(
+			Payment::on_idle(System::block_number(), Weight::MAX),
+			<()>::remove_task()
+		);
+
+		let payment_amount = 20;
+		let expected_cancel_block = CANCEL_BLOCK_BUFFER + 1;
+
+		assert_ok!(Payment::pay(
+			RuntimeOrigin::signed(PAYMENT_CREATOR),
+			PAYMENT_RECIPENT,
+			CURRENCY_ID,
+			payment_amount,
+			None
+		));
+
+		// creator requests a refund
+		assert_ok!(Payment::request_refund(
+			RuntimeOrigin::signed(PAYMENT_CREATOR),
+			PAYMENT_RECIPENT
+		));
+		// ensure the request is added to the refund queue
+		let scheduled_tasks_list = ScheduledTasks::<Test>::get();
+		assert_eq!(scheduled_tasks_list.len(), 1);
+		assert_eq!(
+			scheduled_tasks_list.get(&(PAYMENT_CREATOR, PAYMENT_RECIPENT)).unwrap(),
+			&ScheduledTask {
+				task: Task::Cancel,
+				when: expected_cancel_block
+			}
+		);
+
+		assert_eq!(run_n_blocks(CANCEL_BLOCK_BUFFER - 1), 600);
+		assert_eq!(
+			Payment::on_idle(System::block_number(), Weight::MAX),
+			<()>::remove_task()
+		);
+
+		assert_eq!(run_n_blocks(1), 601);
+		assert_eq!(
+			Payment::on_idle(System::block_number(), Weight::MAX),
+			<()>::remove_task() + <()>::cancel()
+		);
+	});
+}