Cleanup preimage after tlc is settled

src/fiber/channel.rs

@@ -1080,6 +1080,10 @@ where
                 .map_err(|_| {
                     ProcessingChannelError::InternalError("insert preimage failed".to_string())
                 })?;
+            debug_event!(
+                self.network,
+                &format!("store payment_preimage for: {:?}", payment_hash)
+            );
         }
         Ok(())
     }
@@ -1092,12 +1096,15 @@ where
     ) -> Result<(), ProcessingChannelError> {
         let channel_id = state.get_id();
         let (tlc_info, remove_reason) = state.remove_tlc_with_reason(tlc_id)?;
-        if matches!(remove_reason, RemoveTlcReason::RemoveTlcFulfill(_))
-            && self.store.get_invoice(&tlc_info.payment_hash).is_some()
-        {
+        if matches!(remove_reason, RemoveTlcReason::RemoveTlcFulfill(_)) {
+            if self.store.get_invoice(&tlc_info.payment_hash).is_some() {
+                self.store
+                    .update_invoice_status(&tlc_info.payment_hash, CkbInvoiceStatus::Paid)
+                    .expect("update invoice status failed");
+            }
             self.store
-                .update_invoice_status(&tlc_info.payment_hash, CkbInvoiceStatus::Paid)
-                .expect("update invoice status failed");
+                .remove_payment_preimage(&tlc_info.payment_hash)
+                .expect("remove preimage failed");
         }
 
         if let (

src/fiber/tests/channel.rs

@@ -425,16 +425,11 @@ async fn test_network_send_payment_normal_keysend_workflow() {
     assert_eq!(res.status, PaymentSessionStatus::Created);
     let payment_hash = res.payment_hash;
 
-    tokio::time::sleep(tokio::time::Duration::from_millis(2000)).await;
-    let message = |rpc_reply| -> NetworkActorMessage {
-        NetworkActorMessage::Command(NetworkActorCommand::GetPayment(payment_hash, rpc_reply))
-    };
-    let res = call!(node_a.network_actor, message)
-        .expect("node_a alive")
-        .unwrap();
-
-    assert_eq!(res.status, PaymentSessionStatus::Success);
     assert_eq!(res.failed_error, None);
+    node_a.wait_until_success(payment_hash).await;
+
+    let payment_preimage = node_a.get_payment_preimage(&payment_hash);
+    assert!(payment_preimage.is_none());
 }
 
 #[tokio::test]
@@ -1175,12 +1170,13 @@ async fn test_network_send_payment_with_dry_run() {
 async fn test_send_payment_with_3_nodes() {
     init_tracing();
     let _span = tracing::info_span!("node", node = "test").entered();
-    let (node_a, node_b, node_c, channel_1, channel_2) = create_3_nodes_with_established_channel(
-        (100000000000, 100000000000),
-        (100000000000, 100000000000),
-        true,
-    )
-    .await;
+    let (node_a, mut node_b, node_c, channel_1, channel_2) =
+        create_3_nodes_with_established_channel(
+            (100000000000, 100000000000),
+            (100000000000, 100000000000),
+            true,
+        )
+        .await;
     let node_a_local = node_a.get_local_balance_from_channel(channel_1);
     let node_b_local_left = node_b.get_local_balance_from_channel(channel_1);
     let node_b_local_right = node_b.get_local_balance_from_channel(channel_2);
@@ -1216,16 +1212,17 @@ async fn test_send_payment_with_3_nodes() {
     let res = res.unwrap();
     assert_eq!(res.status, PaymentSessionStatus::Created);
     assert!(res.fee > 0);
-    // sleep for 2 seconds to make sure the payment is sent
-    tokio::time::sleep(tokio::time::Duration::from_millis(2000)).await;
-    let message = |rpc_reply| -> NetworkActorMessage {
-        NetworkActorMessage::Command(NetworkActorCommand::GetPayment(res.payment_hash, rpc_reply))
-    };
-    let res = call!(node_a.network_actor, message)
-        .expect("node_a alive")
-        .unwrap();
-    assert_eq!(res.status, PaymentSessionStatus::Success);
-    assert_eq!(res.failed_error, None);
+
+    node_b
+        .expect_debug_event(&format!(
+            "store payment_preimage for: {:?}",
+            res.payment_hash
+        ))
+        .await;
+    assert!(node_b.get_payment_preimage(&res.payment_hash).is_some());
+
+    node_a.wait_until_success(res.payment_hash).await;
+    assert!(node_b.get_payment_preimage(&res.payment_hash).is_none());
 
     let new_node_a_local = node_a.get_local_balance_from_channel(channel_1);
     let new_node_b_left = node_b.get_local_balance_from_channel(channel_1);
@@ -4418,13 +4415,8 @@ async fn test_shutdown_channel_with_different_size_shutdown_script() {
         })
         .await;
 
-    node_a
-        .expect_event(|event| matches!(event, NetworkServiceEvent::DebugEvent(DebugEvent::Common(message)) if message == "ChannelClosed"))
-        .await;
-
-    node_b
-        .expect_event(|event| matches!(event, NetworkServiceEvent::DebugEvent(DebugEvent::Common(message)) if message == "ChannelClosed"))
-        .await;
+    node_a.expect_debug_event("ChannelClosed").await;
+    node_b.expect_debug_event("ChannelClosed").await;
 
     assert_eq!(node_a_shutdown_tx_hash, node_b_shutdown_tx_hash);
 
@@ -5477,6 +5469,9 @@ async fn test_send_payment_will_succeed_with_valid_invoice() {
         node_3.get_invoice_status(ckb_invoice.payment_hash()),
         Some(CkbInvoiceStatus::Paid)
     );
+    assert!(node_3
+        .get_payment_preimage(&ckb_invoice.payment_hash())
+        .is_none());
 }
 
 #[tokio::test]
@@ -5618,6 +5613,9 @@ async fn test_send_payment_will_fail_with_cancelled_invoice() {
         node_3.get_invoice_status(ckb_invoice.payment_hash()),
         Some(CkbInvoiceStatus::Cancelled)
     );
+    assert!(node_3
+        .get_payment_preimage(&ckb_invoice.payment_hash())
+        .is_some());
 }
 
 #[tokio::test]

src/fiber/tests/test_utils.rs

@@ -5,6 +5,7 @@ use crate::fiber::channel::ChannelCommandWithId;
 use crate::fiber::graph::NetworkGraphStateStore;
 use crate::fiber::graph::PaymentSession;
 use crate::fiber::graph::PaymentSessionStatus;
+use crate::fiber::network::DebugEvent;
 use crate::fiber::network::NodeInfoResponse;
 use crate::fiber::network::SendPaymentCommand;
 use crate::fiber::network::SendPaymentResponse;
@@ -543,6 +544,10 @@ impl NetworkNode {
             .expect("cancell success");
     }
 
+    pub fn get_payment_preimage(&self, payment_hash: &Hash256) -> Option<Hash256> {
+        self.store.get_invoice_preimage(payment_hash)
+    }
+
     pub async fn send_payment(
         &mut self,
         command: SendPaymentCommand,
@@ -973,6 +978,13 @@ impl NetworkNode {
             .await;
     }
 
+    pub async fn expect_debug_event(&mut self, message: &str) {
+        self.expect_event(|event| {
+            matches!(event, NetworkServiceEvent::DebugEvent(DebugEvent::Common(msg)) if msg == message)
+        })
+        .await;
+    }
+
     pub async fn submit_tx(&mut self, tx: TransactionView) -> ckb_jsonrpc_types::Status {
         submit_tx(self.chain_actor.clone(), tx).await
     }

src/invoice/store.rs

@@ -20,4 +20,5 @@ pub trait InvoiceStore {
         payment_hash: Hash256,
         preimage: Hash256,
     ) -> Result<(), InvoiceError>;
+    fn remove_payment_preimage(&self, payment_hash: &Hash256) -> Result<(), InvoiceError>;
 }

src/store/store.rs

@@ -512,6 +512,13 @@ impl InvoiceStore for Store {
         batch.commit();
         Ok(())
     }
+
+    fn remove_payment_preimage(&self, payment_hash: &Hash256) -> Result<(), InvoiceError> {
+        let mut batch = self.batch();
+        batch.delete([&[CKB_INVOICE_PREIMAGE_PREFIX], payment_hash.as_ref()].concat());
+        batch.commit();
+        Ok(())
+    }
 }
 
 impl NetworkGraphStateStore for Store {