Update owned channel balance in graph

src/fiber/channel.rs

@@ -479,22 +479,7 @@ where
             }
             FiberChannelMessage::UpdateTlcInfo(update_tlc_info) => {
                 state.remote_tlc_info = Some(update_tlc_info.into());
-                let channel_outpoint = state.must_get_funding_transaction_outpoint();
-                let peer_id = state.get_remote_pubkey();
-                let channel_update_info = state
-                    .get_remote_channel_update_info()
-                    .expect("remote tlc info set above; qed");
-                self.network
-                    .send_message(NetworkActorMessage::new_event(
-                        NetworkActorEvent::OwnedChannelUpdateEvent(
-                            super::graph::OwnedChannelUpdateEvent::Updated(
-                                channel_outpoint,
-                                peer_id,
-                                channel_update_info,
-                            ),
-                        ),
-                    ))
-                    .expect(ASSUME_NETWORK_ACTOR_ALIVE);
+                state.update_graph_for_remote_channel_change(&self.network);
                 Ok(())
             }
             FiberChannelMessage::AddTlc(add_tlc) => {
@@ -711,6 +696,7 @@ where
         state: &mut ChannelActorState,
         inbound: bool,
     ) {
+        let previous_balance = state.get_local_balance();
         let pending_tlcs = if inbound {
             state.tlc_state.received_tlcs.tlcs.iter_mut()
         } else {
@@ -733,6 +719,10 @@ where
                 .await
                 .expect("expect remove tlc success");
         }
+        if state.get_local_balance() != previous_balance {
+            state.update_graph_for_local_channel_change(&self.network);
+            state.update_graph_for_remote_channel_change(&self.network);
+        }
     }
 
     async fn process_add_tlc_error(
@@ -3688,12 +3678,29 @@ impl ChannelActorState {
                 ))
                 .expect(ASSUME_NETWORK_ACTOR_ALIVE);
         }
-        self.update_graph_for_channel_change(network);
+        self.update_graph_for_local_channel_change(network);
         self.send_update_tlc_info_message(network);
     }
 
-    fn update_graph_for_channel_change(&mut self, network: &ActorRef<NetworkActorMessage>) {
-        // Also update network graph with latest local channel update info.
+    fn update_graph_for_remote_channel_change(&mut self, network: &ActorRef<NetworkActorMessage>) {
+        if let Some(channel_update_info) = self.get_remote_channel_update_info() {
+            let channel_outpoint = self.must_get_funding_transaction_outpoint();
+            let peer_id = self.get_remote_pubkey();
+            network
+                .send_message(NetworkActorMessage::new_event(
+                    NetworkActorEvent::OwnedChannelUpdateEvent(
+                        super::graph::OwnedChannelUpdateEvent::Updated(
+                            channel_outpoint,
+                            peer_id,
+                            channel_update_info,
+                        ),
+                    ),
+                ))
+                .expect(ASSUME_NETWORK_ACTOR_ALIVE);
+        }
+    }
+
+    fn update_graph_for_local_channel_change(&mut self, network: &ActorRef<NetworkActorMessage>) {
         let channel_outpoint = self.must_get_funding_transaction_outpoint();
         let peer_id = self.get_local_pubkey();
         let channel_update_info = self.get_local_channel_update_info();

src/fiber/graph.rs

@@ -154,6 +154,17 @@ impl ChannelInfo {
             .map(|n| n.timestamp)
             .max(self.update_of_node1.as_ref().map(|n| n.timestamp))
     }
+
+    #[cfg(test)]
+    pub fn get_channel_update_of(&self, node: Pubkey) -> Option<&ChannelUpdateInfo> {
+        if self.node1() == node {
+            self.update_of_node1.as_ref()
+        } else if self.node2() == node {
+            self.update_of_node2.as_ref()
+        } else {
+            None
+        }
+    }
 }
 
 impl TryFrom<&ChannelActorState> for ChannelInfo {

src/fiber/tests/channel.rs

@@ -1,11 +1,11 @@
 use super::test_utils::{init_tracing, NetworkNode};
 use crate::fiber::channel::UpdateCommand;
 use crate::fiber::config::MAX_PAYMENT_TLC_EXPIRY_LIMIT;
-use crate::fiber::graph::PaymentSessionStatus;
+use crate::fiber::graph::{ChannelInfo, PaymentSessionStatus};
 use crate::fiber::network::{DebugEvent, SendPaymentCommand};
 use crate::fiber::tests::test_utils::*;
 use crate::fiber::types::{
-    Hash256, PaymentHopData, PeeledOnionPacket, TlcErrorCode, NO_SHARED_SECRET,
+    Hash256, PaymentHopData, PeeledOnionPacket, Pubkey, TlcErrorCode, NO_SHARED_SECRET,
 };
 use crate::invoice::{CkbInvoiceStatus, Currency, InvoiceBuilder};
 use crate::{
@@ -280,6 +280,175 @@ async fn test_owned_private_channel_saved_to_the_owner_graph() {
     do_test_owned_channel_saved_to_the_owner_graph(false).await;
 }
 
+#[tokio::test]
+async fn test_update_graph_balance_after_payment() {
+    init_tracing();
+
+    let _span = tracing::info_span!("node", node = "test").entered();
+    let node_a_funding_amount = 100000000000;
+    let node_b_funding_amount = 6200000000;
+
+    let (node_a, node_b, new_channel_id) =
+        create_nodes_with_established_channel(node_a_funding_amount, node_b_funding_amount, true)
+            .await;
+    let node_a_pubkey = node_a.pubkey;
+    let node_b_pubkey = node_b.pubkey;
+
+    // Wait for the channel announcement to be broadcasted
+    tokio::time::sleep(tokio::time::Duration::from_millis(2000)).await;
+
+    let test_channel_info = |channels: Vec<ChannelInfo>,
+                             node_a_pubkey: Pubkey,
+                             node_b_pubkey: Pubkey,
+                             node_a_balance: u128,
+                             node_b_balance: u128| {
+        assert_eq!(channels.len(), 1);
+        let channel = &channels[0];
+        assert_ne!(channel.update_of_node1, None);
+        assert_ne!(channel.update_of_node2, None);
+        assert_ne!(channel.get_channel_update_of(node_a_pubkey), None);
+        assert_ne!(channel.get_channel_update_of(node_b_pubkey), None);
+        assert_eq!(
+            channel
+                .get_channel_update_of(node_a_pubkey)
+                .unwrap()
+                .receivable_balance,
+            Some(node_b_balance)
+        );
+        assert_eq!(
+            channel
+                .get_channel_update_of(node_b_pubkey)
+                .unwrap()
+                .receivable_balance,
+            Some(node_a_balance)
+        );
+    };
+
+    let node_a_old_balance = node_a.get_local_balance_from_channel(new_channel_id);
+    let node_b_old_balance = node_b.get_local_balance_from_channel(new_channel_id);
+    let node_a_old_channels = node_a.get_network_graph_channels().await;
+    let node_b_old_channels = node_b.get_network_graph_channels().await;
+    test_channel_info(
+        node_a_old_channels,
+        node_a_pubkey,
+        node_b_pubkey,
+        node_a_old_balance,
+        node_b_old_balance,
+    );
+    test_channel_info(
+        node_b_old_channels,
+        node_a_pubkey,
+        node_b_pubkey,
+        node_a_old_balance,
+        node_b_old_balance,
+    );
+
+    let message = |rpc_reply| -> NetworkActorMessage {
+        NetworkActorMessage::Command(NetworkActorCommand::SendPayment(
+            SendPaymentCommand {
+                target_pubkey: Some(node_b_pubkey),
+                amount: Some(10000),
+                payment_hash: None,
+                final_tlc_expiry_delta: None,
+                tlc_expiry_limit: None,
+                invoice: None,
+                timeout: None,
+                max_fee_amount: None,
+                max_parts: None,
+                keysend: Some(true),
+                udt_type_script: None,
+                allow_self_payment: false,
+                hop_hints: None,
+                dry_run: false,
+            },
+            rpc_reply,
+        ))
+    };
+    let res1 = call!(node_a.network_actor, message)
+        .expect("node_a alive")
+        .unwrap();
+    // this is the payment_hash generated by keysend
+    assert_eq!(res1.status, PaymentSessionStatus::Created);
+    let payment_hash1 = res1.payment_hash;
+
+    // the second payment is send from node_b to node_a
+    let message = |rpc_reply| -> NetworkActorMessage {
+        NetworkActorMessage::Command(NetworkActorCommand::SendPayment(
+            SendPaymentCommand {
+                target_pubkey: Some(node_a_pubkey),
+                amount: Some(9999),
+                payment_hash: None,
+                final_tlc_expiry_delta: None,
+                invoice: None,
+                timeout: None,
+                tlc_expiry_limit: None,
+                max_fee_amount: None,
+                max_parts: None,
+                keysend: Some(true),
+                udt_type_script: None,
+                allow_self_payment: false,
+                hop_hints: None,
+                dry_run: false,
+            },
+            rpc_reply,
+        ))
+    };
+    let res2 = call!(node_b.network_actor, message)
+        .expect("node_a alive")
+        .unwrap();
+    // this is the payment_hash generated by keysend
+    assert_eq!(res2.status, PaymentSessionStatus::Created);
+    let payment_hash2 = res2.payment_hash;
+
+    // sleep for 2 seconds to make sure the payment is processed
+    tokio::time::sleep(tokio::time::Duration::from_millis(2000)).await;
+
+    let message = |rpc_reply| -> NetworkActorMessage {
+        NetworkActorMessage::Command(NetworkActorCommand::GetPayment(payment_hash1, 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);
+
+    let message = |rpc_reply| -> NetworkActorMessage {
+        NetworkActorMessage::Command(NetworkActorCommand::GetPayment(payment_hash2, rpc_reply))
+    };
+    let res = call!(node_b.network_actor, message)
+        .expect("node_a alive")
+        .unwrap();
+
+    assert_eq!(res.status, PaymentSessionStatus::Success);
+    assert_eq!(res.failed_error, None);
+
+    let node_a_new_balance = node_a.get_local_balance_from_channel(new_channel_id);
+    let node_b_new_balance = node_b.get_local_balance_from_channel(new_channel_id);
+
+    // assert the balance is right,
+    // node_a send 10000 to node_b, and node_b send 9999 to node_a
+    // so the balance should be node_a_old_balance - 1, node_b_old_balance + 1
+    assert_eq!(node_a_new_balance, node_a_old_balance - 1);
+    assert_eq!(node_b_new_balance, node_b_old_balance + 1);
+
+    let node_a_new_channels = node_a.get_network_graph_channels().await;
+    let node_b_new_channels = node_b.get_network_graph_channels().await;
+    test_channel_info(
+        node_a_new_channels,
+        node_a_pubkey,
+        node_b_pubkey,
+        node_a_new_balance,
+        node_b_new_balance,
+    );
+    test_channel_info(
+        node_b_new_channels,
+        node_a_pubkey,
+        node_b_pubkey,
+        node_a_new_balance,
+        node_b_new_balance,
+    );
+}
+
 #[tokio::test]
 async fn test_public_channel_saved_to_the_other_nodes_graph() {
     init_tracing();