f Replace persistence guard with simple update method

src/event.rs

@@ -337,7 +337,7 @@ where
 						);
 						self.channel_manager.fail_htlc_backwards(&payment_hash);
 						self.payment_store
-							.set_status(&payment_hash, PaymentStatus::Failed)
+							.update(&payment_hash, None, None, None, Some(PaymentStatus::Failed))
 							.expect("Failed to access payment store");
 						return;
 					}
@@ -372,7 +372,7 @@ where
 					);
 					self.channel_manager.fail_htlc_backwards(&payment_hash);
 					self.payment_store
-						.set_status(&payment_hash, PaymentStatus::Failed)
+						.update(&payment_hash, None, None, None, Some(PaymentStatus::Failed))
 						.expect("Failed to access payment store");
 				}
 			}
@@ -388,31 +388,68 @@ where
 					hex_utils::to_string(&payment_hash.0),
 					amount_msat,
 				);
-				let (payment_preimage, payment_secret) = match purpose {
+				match purpose {
 					PaymentPurpose::InvoicePayment { payment_preimage, payment_secret, .. } => {
-						(payment_preimage, Some(payment_secret))
+						match self.payment_store.update(
+							&payment_hash,
+							Some(payment_preimage),
+							Some(Some(payment_secret)),
+							Some(Some(amount_msat)),
+							Some(PaymentStatus::Succeeded),
+						) {
+							Ok(true) => (),
+							Ok(false) => {
+								log_error!(
+									self.logger,
+									"Payment with hash {} couldn't be found in store",
+									hex_utils::to_string(&payment_hash.0)
+								);
+								debug_assert!(false);
+							}
+							Err(e) => {
+								log_error!(
+									self.logger,
+									"Failed to update payment with hash {}: {}",
+									hex_utils::to_string(&payment_hash.0),
+									e
+								);
+								debug_assert!(false);
+							}
+						}
+					}
+					PaymentPurpose::SpontaneousPayment(preimage) => {
+						let payment_info = PaymentInfo {
+							preimage: Some(preimage),
+							payment_hash,
+							secret: None,
+							amount_msat: Some(amount_msat),
+							direction: PaymentDirection::Inbound,
+							status: PaymentStatus::Succeeded,
+						};
+
+						match self.payment_store.insert(payment_info) {
+							Ok(false) => (),
+							Ok(true) => {
+								log_error!(
+									self.logger,
+									"Spontaneous payment with hash {} was previosly known",
+									hex_utils::to_string(&payment_hash.0)
+								);
+								debug_assert!(false);
+							}
+							Err(e) => {
+								log_error!(
+									self.logger,
+									"Failed to insert payment with hash {}: {}",
+									hex_utils::to_string(&payment_hash.0),
+									e
+								);
+								debug_assert!(false);
+							}
+						}
 					}
-					PaymentPurpose::SpontaneousPayment(preimage) => (Some(preimage), None),
 				};
 
-				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);
-					})
-					.or_insert(PaymentInfo {
-						preimage: payment_preimage,
-						payment_hash,
-						secret: payment_secret,
-						amount_msat: Some(amount_msat),
-						direction: PaymentDirection::Inbound,
-						status: PaymentStatus::Succeeded,
-					});
-
 				self.event_queue
 					.add_event(Event::PaymentReceived { payment_hash, amount_msat })
 					.expect("Failed to push to event queue");
@@ -450,7 +487,7 @@ where
 				);
 
 				self.payment_store
-					.set_status(&payment_hash, PaymentStatus::Failed)
+					.update(&payment_hash, None, None, None, Some(PaymentStatus::Failed))
 					.expect("Failed to access payment store");
 				self.event_queue
 					.add_event(Event::PaymentFailed { payment_hash })

src/lib.rs

@@ -1042,7 +1042,13 @@ impl Node {
 			Err(payment::PaymentError::Sending(e)) => {
 				log_error!(self.logger, "Failed to send payment: {:?}", e);
 
-				self.payment_store.set_status(&payment_hash, PaymentStatus::Failed)?;
+				self.payment_store.update(
+					&payment_hash,
+					None,
+					None,
+					None,
+					Some(PaymentStatus::Failed),
+				)?;
 				Err(Error::PaymentFailed)
 			}
 		}
@@ -1131,7 +1137,13 @@ impl Node {
 			Err(payment::PaymentError::Sending(e)) => {
 				log_error!(self.logger, "Failed to send payment: {:?}", e);
 
-				self.payment_store.set_status(&payment_hash, PaymentStatus::Failed)?;
+				self.payment_store.update(
+					&payment_hash,
+					None,
+					None,
+					None,
+					Some(PaymentStatus::Failed),
+				)?;
 				Err(Error::PaymentFailed)
 			}
 		}

src/payment_store.rs

@@ -7,11 +7,10 @@ use lightning::ln::{PaymentHash, PaymentPreimage, PaymentSecret};
 use lightning::util::ser::Writeable;
 use lightning::{impl_writeable_tlv_based, impl_writeable_tlv_based_enum};
 
-use std::collections::hash_map;
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use std::iter::FromIterator;
 use std::ops::Deref;
-use std::sync::{Mutex, MutexGuard};
+use std::sync::Mutex;
 
 /// Represents a payment.
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -97,17 +96,13 @@ where
 		Self { payments, kv_store, logger }
 	}
 
-	pub(crate) fn insert(&self, payment_info: PaymentInfo) -> Result<(), Error> {
+	pub(crate) fn insert(&self, payment_info: PaymentInfo) -> Result<bool, Error> {
 		let mut locked_payments = self.payments.lock().unwrap();
 
 		let payment_hash = payment_info.payment_hash.clone();
-		locked_payments.insert(payment_hash.clone(), payment_info.clone());
-		self.write_info_and_commit(&payment_hash, &payment_info)
-	}
-
-	pub(crate) fn lock(&self) -> Result<PaymentInfoGuard<K>, ()> {
-		let locked_store = self.payments.lock().map_err(|_| ())?;
-		Ok(PaymentInfoGuard::new(locked_store, self.kv_store.clone()))
+		let updated = locked_payments.insert(payment_hash.clone(), payment_info.clone()).is_some();
+		self.write_info_and_commit(&payment_hash, &payment_info)?;
+		Ok(updated)
 	}
 
 	pub(crate) fn remove(&self, payment_hash: &PaymentHash) -> Result<(), Error> {
@@ -133,16 +128,36 @@ where
 		self.payments.lock().unwrap().contains_key(payment_hash)
 	}
 
-	pub(crate) fn set_status(
-		&self, payment_hash: &PaymentHash, payment_status: PaymentStatus,
-	) -> Result<(), Error> {
+	pub(crate) fn update(
+		&self, payment_hash: &PaymentHash, update_preimage: Option<Option<PaymentPreimage>>,
+		update_secret: Option<Option<PaymentSecret>>, update_amount_msat: Option<Option<u64>>,
+		update_status: Option<PaymentStatus>,
+	) -> Result<bool, Error> {
+		let mut updated = false;
 		let mut locked_payments = self.payments.lock().unwrap();
 
 		if let Some(payment_info) = locked_payments.get_mut(payment_hash) {
-			payment_info.status = payment_status;
+			if let Some(preimage_opt) = update_preimage {
+				payment_info.preimage = preimage_opt;
+			}
+
+			if let Some(secret_opt) = update_secret {
+				payment_info.secret = secret_opt;
+			}
+
+			if let Some(amount_opt) = update_amount_msat {
+				payment_info.amount_msat = amount_opt;
+			}
+
+			if let Some(status) = update_status {
+				payment_info.status = status;
+			}
+
 			self.write_info_and_commit(payment_hash, payment_info)?;
+			updated = true;
 		}
-		Ok(())
+
+		Ok(updated)
 	}
 
 	fn write_info_and_commit(
@@ -184,59 +199,6 @@ where
 	}
 }
 
-pub(crate) struct PaymentInfoGuard<'a, K: Deref>
-where
-	K::Target: KVStore,
-{
-	inner: MutexGuard<'a, HashMap<PaymentHash, PaymentInfo>>,
-	touched_keys: HashSet<PaymentHash>,
-	kv_store: K,
-}
-
-impl<'a, K: Deref> PaymentInfoGuard<'a, K>
-where
-	K::Target: KVStore,
-{
-	pub fn new(inner: MutexGuard<'a, HashMap<PaymentHash, PaymentInfo>>, kv_store: K) -> Self {
-		let touched_keys = HashSet::new();
-		Self { inner, touched_keys, kv_store }
-	}
-
-	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: KVStore,
-{
-	fn drop(&mut self) {
-		for key in self.touched_keys.iter() {
-			let store_key = hex_utils::to_string(&key.0);
-
-			match self.inner.entry(*key) {
-				hash_map::Entry::Vacant(_) => {
-					self.kv_store
-						.remove(PAYMENT_INFO_PERSISTENCE_NAMESPACE, &store_key)
-						.expect("Persistence failed");
-				}
-				hash_map::Entry::Occupied(e) => {
-					let mut writer = self
-						.kv_store
-						.write(PAYMENT_INFO_PERSISTENCE_NAMESPACE, &store_key)
-						.expect("Persistence failed");
-					e.get().write(&mut writer).expect("Persistence failed");
-					writer.commit().expect("Persistence failed");
-				}
-			};
-		}
-	}
-}
-
 #[cfg(test)]
 mod tests {
 	use super::*;
@@ -262,7 +224,25 @@ mod tests {
 		};
 
 		assert!(!store.get_and_clear_did_persist());
-		payment_info_store.lock().unwrap().entry(payment_hash).or_insert(payment_info);
+
+		assert_eq!(Ok(false), payment_info_store.insert(payment_info.clone()));
 		assert!(store.get_and_clear_did_persist());
+
+		assert_eq!(Ok(true), payment_info_store.insert(payment_info));
+		assert!(store.get_and_clear_did_persist());
+
+		assert_eq!(
+			Ok(true),
+			payment_info_store.update(
+				&payment_hash,
+				None,
+				None,
+				None,
+				Some(PaymentStatus::Succeeded)
+			)
+		);
+		assert!(store.get_and_clear_did_persist());
+
+		assert_eq!(PaymentStatus::Succeeded, payment_info_store.get(&payment_hash).unwrap().status);
 	}
 }