Do DB insertions in parallel
[rapid-gossip-sync-server] / src / persistence.rs
1 use std::fs::OpenOptions;
2 use std::io::{BufWriter, Write};
3 use std::ops::Deref;
4 use std::mem;
5 use std::sync::Arc;
6 use std::time::{Duration, Instant};
7 use lightning::log_info;
8 use lightning::routing::gossip::NetworkGraph;
9 use lightning::util::logger::Logger;
10 use lightning::util::ser::Writeable;
11 use tokio::sync::{mpsc, Mutex, Semaphore};
12
13 use crate::config;
14 use crate::types::GossipMessage;
15
16 const POSTGRES_INSERT_TIMEOUT: Duration = Duration::from_secs(15);
17 const INSERT_PARALELLISM: usize = 16;
18
19 pub(crate) struct GossipPersister<L: Deref> where L::Target: Logger {
20         gossip_persistence_receiver: mpsc::Receiver<GossipMessage>,
21         network_graph: Arc<NetworkGraph<L>>,
22         logger: L
23 }
24
25 impl<L: Deref> GossipPersister<L> where L::Target: Logger {
26         pub fn new(network_graph: Arc<NetworkGraph<L>>, logger: L) -> (Self, mpsc::Sender<GossipMessage>) {
27                 let (gossip_persistence_sender, gossip_persistence_receiver) =
28                         mpsc::channel::<GossipMessage>(100);
29                 (GossipPersister {
30                         gossip_persistence_receiver,
31                         network_graph,
32                         logger
33                 }, gossip_persistence_sender)
34         }
35
36         pub(crate) async fn persist_gossip(&mut self) {
37                 let mut client = crate::connect_to_db().await;
38
39                 {
40                         // initialize the database
41                         let initialization = client
42                                 .execute(config::db_config_table_creation_query(), &[])
43                                 .await;
44                         if let Err(initialization_error) = initialization {
45                                 panic!("db init error: {}", initialization_error);
46                         }
47
48                         let cur_schema = client.query("SELECT db_schema FROM config WHERE id = $1", &[&1]).await.unwrap();
49                         if !cur_schema.is_empty() {
50                                 config::upgrade_db(cur_schema[0].get(0), &mut client).await;
51                         }
52
53                         let preparation = client.execute("set time zone UTC", &[]).await;
54                         if let Err(preparation_error) = preparation {
55                                 panic!("db preparation error: {}", preparation_error);
56                         }
57
58                         let initialization = client
59                                 .execute(
60                                         // TODO: figure out a way to fix the id value without Postgres complaining about
61                                         // its value not being default
62                                         "INSERT INTO config (id, db_schema) VALUES ($1, $2) ON CONFLICT (id) DO NOTHING",
63                                         &[&1, &config::SCHEMA_VERSION]
64                                 ).await;
65                         if let Err(initialization_error) = initialization {
66                                 panic!("db init error: {}", initialization_error);
67                         }
68
69                         let initialization = client
70                                 .execute(config::db_announcement_table_creation_query(), &[])
71                                 .await;
72                         if let Err(initialization_error) = initialization {
73                                 panic!("db init error: {}", initialization_error);
74                         }
75
76                         let initialization = client
77                                 .execute(
78                                         config::db_channel_update_table_creation_query(),
79                                         &[],
80                                 )
81                                 .await;
82                         if let Err(initialization_error) = initialization {
83                                 panic!("db init error: {}", initialization_error);
84                         }
85
86                         let initialization = client
87                                 .batch_execute(config::db_index_creation_query())
88                                 .await;
89                         if let Err(initialization_error) = initialization {
90                                 panic!("db init error: {}", initialization_error);
91                         }
92                 }
93
94                 // print log statement every minute
95                 let mut latest_persistence_log = Instant::now() - Duration::from_secs(60);
96                 let mut i = 0u32;
97                 let mut latest_graph_cache_time = Instant::now();
98                 let insert_limiter = Arc::new(Semaphore::new(INSERT_PARALELLISM));
99                 let connections_cache = Arc::new(Mutex::new(Vec::with_capacity(INSERT_PARALELLISM)));
100                 #[cfg(test)]
101                 let mut tasks_spawned = Vec::new();
102                 // TODO: it would be nice to have some sort of timeout here so after 10 seconds of
103                 // inactivity, some sort of message could be broadcast signaling the activation of request
104                 // processing
105                 while let Some(gossip_message) = self.gossip_persistence_receiver.recv().await {
106                         i += 1; // count the persisted gossip messages
107
108                         if latest_persistence_log.elapsed().as_secs() >= 60 {
109                                 log_info!(self.logger, "Persisting gossip message #{}", i);
110                                 latest_persistence_log = Instant::now();
111                         }
112
113                         // has it been ten minutes? Just cache it
114                         if latest_graph_cache_time.elapsed().as_secs() >= 600 {
115                                 self.persist_network_graph();
116                                 latest_graph_cache_time = Instant::now();
117                         }
118                         insert_limiter.acquire().await.unwrap().forget();
119
120                         let limiter_ref = Arc::clone(&insert_limiter);
121                         let connections_cache_ref = Arc::clone(&connections_cache);
122                         match gossip_message {
123                                 GossipMessage::ChannelAnnouncement(announcement, seen_override) => {
124                                         let scid = announcement.contents.short_channel_id as i64;
125
126                                         // start with the type prefix, which is already known a priori
127                                         let mut announcement_signed = Vec::new();
128                                         announcement.write(&mut announcement_signed).unwrap();
129
130                                         let _task = tokio::spawn(async move {
131                                                 let client;
132                                                 {
133                                                         let mut connections_set = connections_cache_ref.lock().await;
134                                                         if connections_set.is_empty() {
135                                                                 mem::drop(connections_set);
136                                                                 client = crate::connect_to_db().await;
137                                                         } else {
138                                                                 client = connections_set.pop().unwrap();
139                                                         }
140                                                 }
141                                                 if cfg!(test) && seen_override.is_some() {
142                                                         tokio::time::timeout(POSTGRES_INSERT_TIMEOUT, client
143                                                                 .execute("INSERT INTO channel_announcements (\
144                                                                 short_channel_id, \
145                                                                 announcement_signed, \
146                                                                 seen \
147                                                         ) VALUES ($1, $2, TO_TIMESTAMP($3)) ON CONFLICT (short_channel_id) DO NOTHING", &[
148                                                                         &scid,
149                                                                         &announcement_signed,
150                                                                         &(seen_override.unwrap() as f64)
151                                                                 ])).await.unwrap().unwrap();
152                                                 } else {
153                                                         tokio::time::timeout(POSTGRES_INSERT_TIMEOUT, client
154                                                                 .execute("INSERT INTO channel_announcements (\
155                                                                 short_channel_id, \
156                                                                 announcement_signed \
157                                                         ) VALUES ($1, $2) ON CONFLICT (short_channel_id) DO NOTHING", &[
158                                                                         &scid,
159                                                                         &announcement_signed
160                                                                 ])).await.unwrap().unwrap();
161                                                 }
162                                                 let mut connections_set = connections_cache_ref.lock().await;
163                                                 connections_set.push(client);
164                                                 limiter_ref.add_permits(1);
165                                         });
166                                         #[cfg(test)]
167                                         tasks_spawned.push(_task);
168                                 }
169                                 GossipMessage::ChannelUpdate(update, seen_override) => {
170                                         let scid = update.contents.short_channel_id as i64;
171
172                                         let timestamp = update.contents.timestamp as i64;
173
174                                         let direction = (update.contents.flags & 1) == 1;
175                                         let disable = (update.contents.flags & 2) > 0;
176
177                                         let cltv_expiry_delta = update.contents.cltv_expiry_delta as i32;
178                                         let htlc_minimum_msat = update.contents.htlc_minimum_msat as i64;
179                                         let fee_base_msat = update.contents.fee_base_msat as i32;
180                                         let fee_proportional_millionths =
181                                                 update.contents.fee_proportional_millionths as i32;
182                                         let htlc_maximum_msat = update.contents.htlc_maximum_msat as i64;
183
184                                         // start with the type prefix, which is already known a priori
185                                         let mut update_signed = Vec::new();
186                                         update.write(&mut update_signed).unwrap();
187
188                                         let insertion_statement = if cfg!(test) {
189                                                 "INSERT INTO channel_updates (\
190                                                         short_channel_id, \
191                                                         timestamp, \
192                                                         seen, \
193                                                         channel_flags, \
194                                                         direction, \
195                                                         disable, \
196                                                         cltv_expiry_delta, \
197                                                         htlc_minimum_msat, \
198                                                         fee_base_msat, \
199                                                         fee_proportional_millionths, \
200                                                         htlc_maximum_msat, \
201                                                         blob_signed \
202                                                 ) VALUES ($1, $2, TO_TIMESTAMP($3), $4, $5, $6, $7, $8, $9, $10, $11, $12)  ON CONFLICT DO NOTHING"
203                                         } else {
204                                                 "INSERT INTO channel_updates (\
205                                                         short_channel_id, \
206                                                         timestamp, \
207                                                         channel_flags, \
208                                                         direction, \
209                                                         disable, \
210                                                         cltv_expiry_delta, \
211                                                         htlc_minimum_msat, \
212                                                         fee_base_msat, \
213                                                         fee_proportional_millionths, \
214                                                         htlc_maximum_msat, \
215                                                         blob_signed \
216                                                 ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)  ON CONFLICT DO NOTHING"
217                                         };
218
219                                         // this may not be used outside test cfg
220                                         let _seen_timestamp = seen_override.unwrap_or(timestamp as u32) as f64;
221
222                                         let _task = tokio::spawn(async move {
223                                                 let client;
224                                                 {
225                                                         let mut connections_set = connections_cache_ref.lock().await;
226                                                         if connections_set.is_empty() {
227                                                                 mem::drop(connections_set);
228                                                                 client = crate::connect_to_db().await;
229                                                         } else {
230                                                                 client = connections_set.pop().unwrap();
231                                                         }
232                                                 }
233                                                 tokio::time::timeout(POSTGRES_INSERT_TIMEOUT, client
234                                                         .execute(insertion_statement, &[
235                                                                 &scid,
236                                                                 &timestamp,
237                                                                 #[cfg(test)]
238                                                                         &_seen_timestamp,
239                                                                 &(update.contents.flags as i16),
240                                                                 &direction,
241                                                                 &disable,
242                                                                 &cltv_expiry_delta,
243                                                                 &htlc_minimum_msat,
244                                                                 &fee_base_msat,
245                                                                 &fee_proportional_millionths,
246                                                                 &htlc_maximum_msat,
247                                                                 &update_signed
248                                                         ])).await.unwrap().unwrap();
249                                                 let mut connections_set = connections_cache_ref.lock().await;
250                                                 connections_set.push(client);
251                                                 limiter_ref.add_permits(1);
252                                         });
253                                         #[cfg(test)]
254                                         tasks_spawned.push(_task);
255                                 }
256                         }
257                 }
258                 #[cfg(test)]
259                 for task in tasks_spawned {
260                         task.await;
261                 }
262         }
263
264         fn persist_network_graph(&self) {
265                 log_info!(self.logger, "Caching network graph…");
266                 let cache_path = config::network_graph_cache_path();
267                 let file = OpenOptions::new()
268                         .create(true)
269                         .write(true)
270                         .truncate(true)
271                         .open(&cache_path)
272                         .unwrap();
273                 self.network_graph.remove_stale_channels_and_tracking();
274                 let mut writer = BufWriter::new(file);
275                 self.network_graph.write(&mut writer).unwrap();
276                 writer.flush().unwrap();
277                 log_info!(self.logger, "Cached network graph!");
278         }
279 }