Block tasks if the DB writes get behind, rather than growing the queue
[rapid-gossip-sync-server] / src / downloader.rs
1 use std::sync::{Arc, RwLock};
2
3 use bitcoin::secp256k1::PublicKey;
4 use lightning::ln::msgs::{ChannelAnnouncement, ChannelUpdate, Init, LightningError, NodeAnnouncement, QueryChannelRange, QueryShortChannelIds, ReplyChannelRange, ReplyShortChannelIdsEnd, RoutingMessageHandler};
5 use lightning::routing::gossip::{NetworkGraph, P2PGossipSync};
6 use lightning::util::events::{MessageSendEvent, MessageSendEventsProvider};
7 use tokio::sync::mpsc;
8 use tokio::sync::mpsc::error::TrySendError;
9
10 use crate::{GossipChainAccess, TestLogger};
11 use crate::types::GossipMessage;
12
13 pub(crate) struct GossipCounter {
14         pub(crate) channel_announcements: u64,
15         pub(crate) channel_updates: u64,
16         pub(crate) channel_updates_without_htlc_max_msats: u64,
17         pub(crate) channel_announcements_with_mismatched_scripts: u64
18 }
19
20 impl GossipCounter {
21         pub(crate) fn new() -> Self {
22                 Self {
23                         channel_announcements: 0,
24                         channel_updates: 0,
25                         channel_updates_without_htlc_max_msats: 0,
26                         channel_announcements_with_mismatched_scripts: 0,
27                 }
28         }
29 }
30
31 pub(crate) struct GossipRouter {
32         pub(crate) native_router: Arc<P2PGossipSync<Arc<NetworkGraph<Arc<TestLogger>>>, GossipChainAccess, Arc<TestLogger>>>,
33         pub(crate) counter: RwLock<GossipCounter>,
34         pub(crate) sender: mpsc::Sender<GossipMessage>,
35 }
36
37 impl MessageSendEventsProvider for GossipRouter {
38         fn get_and_clear_pending_msg_events(&self) -> Vec<MessageSendEvent> {
39                 self.native_router.get_and_clear_pending_msg_events()
40         }
41 }
42
43 impl RoutingMessageHandler for GossipRouter {
44         fn handle_node_announcement(&self, msg: &NodeAnnouncement) -> Result<bool, LightningError> {
45                 self.native_router.handle_node_announcement(msg)
46         }
47
48         fn handle_channel_announcement(&self, msg: &ChannelAnnouncement) -> Result<bool, LightningError> {
49                 let mut counter = self.counter.write().unwrap();
50
51                 let output_value = self.native_router.handle_channel_announcement(msg).map_err(|error| {
52                         let error_string = format!("{:?}", error);
53                         if error_string.contains("announced on an unknown chain"){
54                                 return error;
55                         }
56                         counter.channel_announcements_with_mismatched_scripts += 1;
57                         error
58                 })?;
59
60                 counter.channel_announcements += 1;
61                 let gossip_message = GossipMessage::ChannelAnnouncement(msg.clone());
62                 if let Err(err) = self.sender.try_send(gossip_message) {
63                         let gossip_message = match err { TrySendError::Full(msg)|TrySendError::Closed(msg) => msg };
64                         tokio::task::block_in_place(move || { tokio::runtime::Handle::current().block_on(async move {
65                                 self.sender.send(gossip_message).await.unwrap();
66                         })});
67                 }
68
69                 Ok(output_value)
70         }
71
72         fn handle_channel_update(&self, msg: &ChannelUpdate) -> Result<bool, LightningError> {
73                 let output_value = self.native_router.handle_channel_update(msg)?;
74
75                 let mut counter = self.counter.write().unwrap();
76                 counter.channel_updates += 1;
77                 let gossip_message = GossipMessage::ChannelUpdate(msg.clone());
78
79                 if let Err(err) = self.sender.try_send(gossip_message) {
80                         let gossip_message = match err { TrySendError::Full(msg)|TrySendError::Closed(msg) => msg };
81                         tokio::task::block_in_place(move || { tokio::runtime::Handle::current().block_on(async move {
82                                 self.sender.send(gossip_message).await.unwrap();
83                         })});
84                 }
85
86                 Ok(output_value)
87         }
88
89         fn get_next_channel_announcements(&self, starting_point: u64, batch_amount: u8) -> Vec<(ChannelAnnouncement, Option<ChannelUpdate>, Option<ChannelUpdate>)> {
90                 self.native_router.get_next_channel_announcements(starting_point, batch_amount)
91         }
92
93         fn get_next_node_announcements(&self, starting_point: Option<&PublicKey>, batch_amount: u8) -> Vec<NodeAnnouncement> {
94                 self.native_router.get_next_node_announcements(starting_point, batch_amount)
95         }
96
97         fn peer_connected(&self, their_node_id: &PublicKey, init: &Init) {
98                 self.native_router.peer_connected(their_node_id, init)
99         }
100
101         fn handle_reply_channel_range(&self, their_node_id: &PublicKey, msg: ReplyChannelRange) -> Result<(), LightningError> {
102                 self.native_router.handle_reply_channel_range(their_node_id, msg)
103         }
104
105         fn handle_reply_short_channel_ids_end(&self, their_node_id: &PublicKey, msg: ReplyShortChannelIdsEnd) -> Result<(), LightningError> {
106                 self.native_router.handle_reply_short_channel_ids_end(their_node_id, msg)
107         }
108
109         fn handle_query_channel_range(&self, their_node_id: &PublicKey, msg: QueryChannelRange) -> Result<(), LightningError> {
110                 self.native_router.handle_query_channel_range(their_node_id, msg)
111         }
112
113         fn handle_query_short_channel_ids(&self, their_node_id: &PublicKey, msg: QueryShortChannelIds) -> Result<(), LightningError> {
114                 self.native_router.handle_query_short_channel_ids(their_node_id, msg)
115         }
116 }