f Expose PaymentInfoStorage::lock and impl ..Guard

Author: tnull

src/event.rs

@@ -182,7 +182,7 @@ impl Writeable for EventQueueSerWrapper<'_> {
 	}
 }
 
-pub(crate) struct EventHandler<K: Deref, L: Deref>
+pub(crate) struct EventHandler<K: Deref + Clone, L: Deref>
 where
 	K::Target: KVStorePersister + KVStoreUnpersister,
 	L::Target: Logger,
@@ -198,7 +198,7 @@ where
 	_config: Arc<Config>,
 }
 
-impl<K: Deref, L: Deref> EventHandler<K, L>
+impl<K: Deref + Clone, L: Deref> EventHandler<K, L>
 where
 	K::Target: KVStorePersister + KVStoreUnpersister,
 	L::Target: Logger,
@@ -223,7 +223,7 @@ where
 	}
 }
 
-impl<K: Deref, L: Deref> LdkEventHandler for EventHandler<K, L>
+impl<K: Deref + Clone, L: Deref> LdkEventHandler for EventHandler<K, L>
 where
 	K::Target: KVStorePersister + KVStoreUnpersister,
 	L::Target: Logger,
@@ -364,25 +364,24 @@ where
 					PaymentPurpose::SpontaneousPayment(preimage) => (Some(preimage), None),
 				};
 
-				let payment_info =
-					if let Some(mut payment_info) = self.payment_store.get(&payment_hash) {
+				let mut locked_store = self.payment_store.lock().unwrap();
+				locked_store
+					.entry(payment_hash)
+					.and_modify(|payment_info| {
 						payment_info.status = PaymentStatus::Succeeded;
 						payment_info.preimage = payment_preimage;
 						payment_info.secret = payment_secret;
 						payment_info.amount_msat = Some(amount_msat);
-						payment_info
-					} else {
-						PaymentInfo {
-							preimage: payment_preimage,
-							payment_hash,
-							secret: payment_secret,
-							amount_msat: Some(amount_msat),
-							direction: PaymentDirection::Inbound,
-							status: PaymentStatus::Succeeded,
-						}
-					};
+					})
+					.or_insert(PaymentInfo {
+						preimage: payment_preimage,
+						payment_hash,
+						secret: payment_secret,
+						amount_msat: Some(amount_msat),
+						direction: PaymentDirection::Inbound,
+						status: PaymentStatus::Succeeded,
+					});
 
-				self.payment_store.insert(payment_info).expect("Failed to access payment store");
 				self.event_queue
 					.add_event(Event::PaymentReceived { payment_hash, amount_msat })
 					.expect("Failed to push to event queue");

src/payment_store.rs

@@ -6,10 +6,11 @@ use lightning::ln::{PaymentHash, PaymentPreimage, PaymentSecret};
 use lightning::util::persist::KVStorePersister;
 use lightning::{impl_writeable_tlv_based, impl_writeable_tlv_based_enum};
 
-use std::collections::HashMap;
+use std::collections::hash_map;
+use std::collections::{HashMap, HashSet};
 use std::iter::FromIterator;
 use std::ops::Deref;
-use std::sync::Mutex;
+use std::sync::{Mutex, MutexGuard};
 
 /// Represents a payment.
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -71,18 +72,23 @@ impl_writeable_tlv_based_enum!(PaymentStatus,
 /// The payment information will be persisted under this prefix.
 pub(crate) const PAYMENT_INFO_PERSISTENCE_PREFIX: &str = "payments";
 
-pub(crate) struct PaymentInfoStorage<K: Deref>
+pub(crate) struct PaymentInfoStorage<K: Deref + Clone>
 where
 	K::Target: KVStorePersister + KVStoreUnpersister,
 {
 	payments: Mutex<HashMap<PaymentHash, PaymentInfo>>,
 	persister: K,
 }
 
-impl<K: Deref> PaymentInfoStorage<K>
+impl<K: Deref + Clone> PaymentInfoStorage<K>
 where
 	K::Target: KVStorePersister + KVStoreUnpersister,
 {
+	pub(crate) fn new(persister: K) -> Self {
+		let payments = Mutex::new(HashMap::new());
+		Self { payments, persister }
+	}
+
 	pub(crate) fn from_payments(payments: Vec<PaymentInfo>, persister: K) -> Self {
 		let payments = Mutex::new(HashMap::from_iter(
 			payments.into_iter().map(|payment_info| (payment_info.payment_hash, payment_info)),
@@ -107,6 +113,11 @@ where
 		return Ok(());
 	}
 
+	pub(crate) fn lock(&self) -> Result<PaymentInfoGuard<K>, ()> {
+		let locked_store = self.payments.lock().map_err(|_| ())?;
+		Ok(PaymentInfoGuard::new(locked_store, self.persister.clone()))
+	}
+
 	pub(crate) fn remove(&self, payment_hash: &PaymentHash) -> Result<(), Error> {
 		let key = format!(
 			"{}/{}",
@@ -143,3 +154,79 @@ where
 		Ok(())
 	}
 }
+
+pub(crate) struct PaymentInfoGuard<'a, K: Deref>
+where
+	K::Target: KVStorePersister + KVStoreUnpersister,
+{
+	inner: MutexGuard<'a, HashMap<PaymentHash, PaymentInfo>>,
+	touched_keys: HashSet<PaymentHash>,
+	persister: K,
+}
+
+impl<'a, K: Deref> PaymentInfoGuard<'a, K>
+where
+	K::Target: KVStorePersister + KVStoreUnpersister,
+{
+	pub fn new(inner: MutexGuard<'a, HashMap<PaymentHash, PaymentInfo>>, persister: K) -> Self {
+		let touched_keys = HashSet::new();
+		Self { inner, touched_keys, persister }
+	}
+
+	pub fn entry(
+		&mut self, payment_hash: PaymentHash,
+	) -> hash_map::Entry<PaymentHash, PaymentInfo> {
+		self.touched_keys.insert(payment_hash);
+		self.inner.entry(payment_hash)
+	}
+}
+
+impl<'a, K: Deref> Drop for PaymentInfoGuard<'a, K>
+where
+	K::Target: KVStorePersister + KVStoreUnpersister,
+{
+	fn drop(&mut self) {
+		for key in self.touched_keys.iter() {
+			let store_key =
+				format!("{}/{}", PAYMENT_INFO_PERSISTENCE_PREFIX, hex_utils::to_string(&key.0));
+
+			match self.inner.entry(*key) {
+				hash_map::Entry::Vacant(_) => {
+					self.persister.unpersist(&store_key).expect("Persistence failed");
+				}
+				hash_map::Entry::Occupied(e) => {
+					self.persister.persist(&store_key, e.get()).expect("Persistence failed");
+				}
+			};
+		}
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use crate::tests::test_utils::TestPersister;
+	use std::sync::Arc;
+
+	#[test]
+	fn persistence_guard_persists_on_drop() {
+		let persister = Arc::new(TestPersister::new());
+		let payment_info_store = PaymentInfoStorage::new(Arc::clone(&persister));
+
+		let payment_hash = PaymentHash([42u8; 32]);
+		assert!(!payment_info_store.contains(&payment_hash));
+
+		let payment_info = PaymentInfo {
+			payment_hash,
+			preimage: None,
+			secret: None,
+			amount_msat: None,
+			direction: PaymentDirection::Inbound,
+			status: PaymentStatus::Pending,
+		};
+
+		assert!(!persister.get_and_clear_did_persist());
+		payment_info_store.lock().unwrap().entry(payment_hash).or_insert(payment_info);
+		assert!(persister.get_and_clear_did_persist());
+	}
+}

src/tests/test_utils.rs

@@ -1,3 +1,4 @@
+use crate::io_utils::KVStoreUnpersister;
 use lightning::util::persist::KVStorePersister;
 use lightning::util::ser::Writeable;
 
@@ -53,3 +54,11 @@ impl KVStorePersister for TestPersister {
 		Ok(())
 	}
 }
+
+impl KVStoreUnpersister for TestPersister {
+	fn unpersist(&self, key: &str) -> std::io::Result<bool> {
+		let mut persisted_bytes_lock = self.persisted_bytes.lock().unwrap();
+		self.did_persist.store(true, Ordering::SeqCst);
+		Ok(persisted_bytes_lock.remove(key).is_some())
+	}
+}