Cleanup network graph persisting and ensure write succeeds
[rapid-gossip-sync-server] / src / persistence.rs
1 use std::fs::OpenOptions;
2 use std::io::{BufWriter, Write};
3 use std::sync::Arc;
4 use std::time::Instant;
5 use lightning::routing::gossip::NetworkGraph;
6 use lightning::util::ser::Writeable;
7 use tokio::sync::mpsc;
8 use tokio_postgres::NoTls;
9
10 use crate::{config, hex_utils, TestLogger};
11 use crate::types::GossipMessage;
12
13 pub(crate) struct GossipPersister {
14         gossip_persistence_receiver: mpsc::Receiver<GossipMessage>,
15         server_sync_completion_sender: mpsc::Sender<()>,
16         network_graph: Arc<NetworkGraph<TestLogger>>,
17 }
18
19 impl GossipPersister {
20         pub fn new(server_sync_completion_sender: mpsc::Sender<()>, network_graph: Arc<NetworkGraph<TestLogger>>) -> (Self, mpsc::Sender<GossipMessage>) {
21                 let (gossip_persistence_sender, gossip_persistence_receiver) =
22                         mpsc::channel::<GossipMessage>(100);
23                 (GossipPersister {
24                         gossip_persistence_receiver,
25                         server_sync_completion_sender,
26                         network_graph
27                 }, gossip_persistence_sender)
28         }
29
30         pub(crate) async fn persist_gossip(&mut self) {
31                 let connection_config = config::db_connection_config();
32                 let (client, connection) =
33                         connection_config.connect(NoTls).await.unwrap();
34
35                 tokio::spawn(async move {
36                         if let Err(e) = connection.await {
37                                 panic!("connection error: {}", e);
38                         }
39                 });
40
41                 {
42                         // initialize the database
43                         let initialization = client
44                                 .execute(config::db_config_table_creation_query(), &[])
45                                 .await;
46                         if let Err(initialization_error) = initialization {
47                                 panic!("db init error: {}", initialization_error);
48                         }
49
50                         let initialization = client
51                                 .execute(
52                                         // TODO: figure out a way to fix the id value without Postgres complaining about
53                                         // its value not being default
54                                         "INSERT INTO config (id, db_schema) VALUES ($1, $2) ON CONFLICT (id) DO NOTHING",
55                                         &[&1, &config::SCHEMA_VERSION]
56                                 ).await;
57                         if let Err(initialization_error) = initialization {
58                                 panic!("db init error: {}", initialization_error);
59                         }
60
61                         let initialization = client
62                                 .execute(config::db_announcement_table_creation_query(), &[])
63                                 .await;
64                         if let Err(initialization_error) = initialization {
65                                 panic!("db init error: {}", initialization_error);
66                         }
67
68                         let initialization = client
69                                 .execute(
70                                         config::db_channel_update_table_creation_query(),
71                                         &[],
72                                 )
73                                 .await;
74                         if let Err(initialization_error) = initialization {
75                                 panic!("db init error: {}", initialization_error);
76                         }
77
78                         let initialization = client
79                                 .batch_execute(config::db_index_creation_query())
80                                 .await;
81                         if let Err(initialization_error) = initialization {
82                                 panic!("db init error: {}", initialization_error);
83                         }
84                 }
85
86                 // print log statement every 10,000 messages
87                 let mut persistence_log_threshold = 10000;
88                 let mut i = 0u32;
89                 let mut server_sync_completion_sent = false;
90                 let mut latest_graph_cache_time = Instant::now();
91                 // TODO: it would be nice to have some sort of timeout here so after 10 seconds of
92                 // inactivity, some sort of message could be broadcast signaling the activation of request
93                 // processing
94                 while let Some(gossip_message) = &self.gossip_persistence_receiver.recv().await {
95                         i += 1; // count the persisted gossip messages
96
97                         if i == 1 || i % persistence_log_threshold == 0 {
98                                 println!("Persisting gossip message #{}", i);
99                         }
100
101                         // has it been ten minutes? Just cache it
102                         if latest_graph_cache_time.elapsed().as_secs() >= 600 {
103                                 self.persist_network_graph();
104                                 latest_graph_cache_time = Instant::now();
105                         }
106
107                         match &gossip_message {
108                                 GossipMessage::InitialSyncComplete => {
109                                         // signal to the server that it may now serve dynamic responses and calculate
110                                         // snapshots
111                                         // we take this detour through the persister to ensure that all previous
112                                         // messages have already been persisted to the database
113                                         println!("Persister caught up with gossip!");
114                                         i -= 1; // this wasn't an actual gossip message that needed persisting
115                                         persistence_log_threshold = 50;
116                                         if !server_sync_completion_sent {
117                                                 server_sync_completion_sent = true;
118                                                 self.server_sync_completion_sender.send(()).await.unwrap();
119                                                 println!("Server has been notified of persistence completion.");
120                                         }
121                                 }
122                                 GossipMessage::ChannelAnnouncement(announcement) => {
123
124                                         let scid = announcement.contents.short_channel_id;
125                                         let scid_hex = hex_utils::hex_str(&scid.to_be_bytes());
126                                         // scid is 8 bytes
127                                         // block height is the first three bytes
128                                         // to obtain block height, shift scid right by 5 bytes (40 bits)
129                                         let block_height = (scid >> 5 * 8) as i32;
130                                         let chain_hash = announcement.contents.chain_hash.as_ref();
131                                         let chain_hash_hex = hex_utils::hex_str(chain_hash);
132
133                                         // start with the type prefix, which is already known a priori
134                                         let mut announcement_signed = Vec::new(); // vec![1, 0];
135                                         announcement.write(&mut announcement_signed).unwrap();
136
137                                         let result = client
138                                                 .execute("INSERT INTO channel_announcements (\
139                                                         short_channel_id, \
140                                                         block_height, \
141                                                         chain_hash, \
142                                                         announcement_signed \
143                                                 ) VALUES ($1, $2, $3, $4) ON CONFLICT (short_channel_id) DO NOTHING", &[
144                                                         &scid_hex,
145                                                         &block_height,
146                                                         &chain_hash_hex,
147                                                         &announcement_signed
148                                                 ]).await;
149                                         if result.is_err() {
150                                                 panic!("error: {}", result.err().unwrap());
151                                         }
152                                 }
153                                 GossipMessage::ChannelUpdate(update) => {
154                                         let scid = update.contents.short_channel_id;
155                                         let scid_hex = hex_utils::hex_str(&scid.to_be_bytes());
156
157                                         let chain_hash = update.contents.chain_hash.as_ref();
158                                         let chain_hash_hex = hex_utils::hex_str(chain_hash);
159
160                                         let timestamp = update.contents.timestamp as i64;
161
162                                         let channel_flags = update.contents.flags as i32;
163                                         let direction = channel_flags & 1;
164                                         let disable = (channel_flags & 2) > 0;
165
166                                         let composite_index = format!("{}:{}:{}", scid_hex, timestamp, direction);
167
168                                         let cltv_expiry_delta = update.contents.cltv_expiry_delta as i32;
169                                         let htlc_minimum_msat = update.contents.htlc_minimum_msat as i64;
170                                         let fee_base_msat = update.contents.fee_base_msat as i32;
171                                         let fee_proportional_millionths =
172                                                 update.contents.fee_proportional_millionths as i32;
173                                         let htlc_maximum_msat = update.contents.htlc_maximum_msat as i64;
174
175                                         // start with the type prefix, which is already known a priori
176                                         let mut update_signed = Vec::new(); // vec![1, 2];
177                                         update.write(&mut update_signed).unwrap();
178
179                                         let result = client
180                                                 .execute("INSERT INTO channel_updates (\
181                                                         composite_index, \
182                                                         chain_hash, \
183                                                         short_channel_id, \
184                                                         timestamp, \
185                                                         channel_flags, \
186                                                         direction, \
187                                                         disable, \
188                                                         cltv_expiry_delta, \
189                                                         htlc_minimum_msat, \
190                                                         fee_base_msat, \
191                                                         fee_proportional_millionths, \
192                                                         htlc_maximum_msat, \
193                                                         blob_signed \
194                                                 ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)  ON CONFLICT (composite_index) DO NOTHING", &[
195                                                         &composite_index,
196                                                         &chain_hash_hex,
197                                                         &scid_hex,
198                                                         &timestamp,
199                                                         &channel_flags,
200                                                         &direction,
201                                                         &disable,
202                                                         &cltv_expiry_delta,
203                                                         &htlc_minimum_msat,
204                                                         &fee_base_msat,
205                                                         &fee_proportional_millionths,
206                                                         &htlc_maximum_msat,
207                                                         &update_signed
208                                                 ]).await;
209                                         if result.is_err() {
210                                                 panic!("error: {}", result.err().unwrap());
211                                         }
212                                 }
213                         }
214                 }
215         }
216
217         fn persist_network_graph(&self) {
218                 println!("Caching network graph…");
219                 let cache_path = config::network_graph_cache_path();
220                 let file = OpenOptions::new()
221                         .create(true)
222                         .write(true)
223                         .truncate(true)
224                         .open(&cache_path)
225                         .unwrap();
226                 self.network_graph.remove_stale_channels();
227                 let mut writer = BufWriter::new(file);
228                 self.network_graph.write(&mut writer).unwrap();
229                 writer.flush().unwrap();
230                 println!("Cached network graph!");
231         }
232 }