f - Make funding-payment classification write atomic with its existence check

Co-Authored-By: Claude <noreply@anthropic.com>

Author: jkczyz

src/data_store.rs

@@ -40,6 +40,13 @@ pub(crate) enum DataStoreUpdateResult {
 	NotFound,
 }
 
+#[derive(PartialEq, Eq, Debug, Clone, Copy)]
+pub(crate) enum DataStoreUpdateOrInsertResult {
+	Inserted,
+	Updated,
+	Unchanged,
+}
+
 pub(crate) struct DataStore<SO: StorableObject, L: Deref>
 where
 	L::Target: LdkLogger,
@@ -81,6 +88,11 @@ where
 		Ok(updated)
 	}
 
+	/// Like [`Self::insert`], but when an entry with the object's id already exists, merges the
+	/// object's full update ([`StorableObject::to_update`]) into it instead of replacing it.
+	///
+	/// Unlike [`Self::update_or_insert`], the caller does not choose what is merged into an
+	/// existing entry: the full update is always applied.
 	pub(crate) async fn insert_or_update(&self, object: SO) -> Result<bool, Error> {
 		let _guard = self.mutation_lock.lock().await;
 
@@ -170,6 +182,47 @@ where
 		Ok(DataStoreUpdateResult::Updated)
 	}
 
+	/// Applies `update` when an object with its id already exists, or inserts `object` when none
+	/// does.
+	///
+	/// Like [`Self::update`], but falls back to inserting `object` instead of returning
+	/// [`DataStoreUpdateResult::NotFound`]. Unlike [`Self::insert_or_update`], the caller chooses
+	/// exactly what is merged into an existing entry: `update` may carry less than the full
+	/// object.
+	///
+	/// The existence check and the write share one critical section of the mutation lock, so a
+	/// concurrent writer cannot land in between and later have its state clobbered by the insert
+	/// fallback — the check-then-act race that separate [`Self::contains_key`] +
+	/// [`Self::insert_or_update`] calls reintroduce.
+	pub(crate) async fn update_or_insert(
+		&self, update: SO::Update, object: SO,
+	) -> Result<DataStoreUpdateOrInsertResult, Error> {
+		debug_assert!(update.id() == object.id(), "update and object must share an id");
+		let _guard = self.mutation_lock.lock().await;
+
+		let id = update.id();
+		let (data_to_persist, result) = {
+			let locked_objects = self.objects.lock().expect("lock");
+			match locked_objects.get(&id) {
+				Some(existing_object) => {
+					let mut updated_object = existing_object.clone();
+					if updated_object.update(update) {
+						(Some(updated_object), DataStoreUpdateOrInsertResult::Updated)
+					} else {
+						(None, DataStoreUpdateOrInsertResult::Unchanged)
+					}
+				},
+				None => (Some(object), DataStoreUpdateOrInsertResult::Inserted),
+			}
+		};
+
+		if let Some(object) = data_to_persist {
+			self.persist(&object).await?;
+			self.objects.lock().expect("lock").insert(id, object);
+		}
+		Ok(result)
+	}
+
 	/// Returns in-memory objects matching `f`.
 	///
 	/// The async mutation lock serializes writers, but this synchronous reader cannot wait on it.
@@ -403,6 +456,102 @@ mod tests {
 		assert_eq!(Ok(true), data_store.insert_or_update(new_iou_object).await);
 	}
 
+	#[tokio::test]
+	async fn update_or_insert_inserts_when_absent() {
+		let store: Arc<DynStore> = Arc::new(DynStoreWrapper(InMemoryStore::new()));
+		let logger = Arc::new(TestLogger::new());
+		let primary_namespace = "datastore_test_primary".to_string();
+		let secondary_namespace = "datastore_test_secondary".to_string();
+		let data_store: DataStore<TestObject, Arc<TestLogger>> = DataStore::new(
+			Vec::new(),
+			primary_namespace.clone(),
+			secondary_namespace.clone(),
+			Arc::clone(&store),
+			logger,
+		);
+
+		let id = TestObjectId { id: [42u8; 4] };
+		let object = TestObject { id, data: [23u8; 3] };
+		let update = TestObjectUpdate { id, data: [25u8; 3] };
+		assert_eq!(
+			Ok(DataStoreUpdateOrInsertResult::Inserted),
+			data_store.update_or_insert(update, object).await
+		);
+
+		// The insert path stores the fallback object as-is; the update is not applied to it.
+		assert_eq!(Some(object), data_store.get(&id));
+		let store_key = id.encode_to_hex_str();
+		assert!(KVStore::read(&*store, &primary_namespace, &secondary_namespace, &store_key)
+			.await
+			.is_ok());
+	}
+
+	#[tokio::test]
+	async fn update_or_insert_applies_update_when_present() {
+		let store: Arc<DynStore> = Arc::new(DynStoreWrapper(InMemoryStore::new()));
+		let logger = Arc::new(TestLogger::new());
+		let id = TestObjectId { id: [42u8; 4] };
+		let existing_object = TestObject { id, data: [23u8; 3] };
+		let data_store: DataStore<TestObject, Arc<TestLogger>> = DataStore::new(
+			vec![existing_object],
+			"datastore_test_primary".to_string(),
+			"datastore_test_secondary".to_string(),
+			store,
+			logger,
+		);
+
+		// When an entry exists, only the update is applied; the fallback object must not replace
+		// it.
+		let update = TestObjectUpdate { id, data: [24u8; 3] };
+		let object = TestObject { id, data: [99u8; 3] };
+		assert_eq!(
+			Ok(DataStoreUpdateOrInsertResult::Updated),
+			data_store.update_or_insert(update, object).await
+		);
+		assert_eq!(data_store.get(&id).unwrap().data, [24u8; 3]);
+	}
+
+	#[tokio::test]
+	async fn update_or_insert_returns_unchanged_without_persisting() {
+		let id = TestObjectId { id: [42u8; 4] };
+		let existing_object = TestObject { id, data: [23u8; 3] };
+		let data_store = new_failing_data_store(vec![existing_object]);
+
+		// A no-op update returns `Unchanged` without attempting to persist (the store fails all
+		// writes) and without falling back to the object.
+		let update = TestObjectUpdate { id, data: [23u8; 3] };
+		let object = TestObject { id, data: [99u8; 3] };
+		assert_eq!(
+			Ok(DataStoreUpdateOrInsertResult::Unchanged),
+			data_store.update_or_insert(update, object).await
+		);
+		assert_eq!(Some(existing_object), data_store.get(&id));
+	}
+
+	#[tokio::test]
+	async fn update_or_insert_does_not_mutate_memory_if_persist_fails() {
+		let existing_id = TestObjectId { id: [42u8; 4] };
+		let existing_object = TestObject { id: existing_id, data: [23u8; 3] };
+		let data_store = new_failing_data_store(vec![existing_object]);
+
+		let update = TestObjectUpdate { id: existing_id, data: [24u8; 3] };
+		let object = TestObject { id: existing_id, data: [24u8; 3] };
+		assert_eq!(
+			Err(Error::PersistenceFailed),
+			data_store.update_or_insert(update, object).await
+		);
+		assert_eq!(Some(existing_object), data_store.get(&existing_id));
+
+		let new_id = TestObjectId { id: [55u8; 4] };
+		let new_object = TestObject { id: new_id, data: [34u8; 3] };
+		let new_update = TestObjectUpdate { id: new_id, data: [34u8; 3] };
+		assert_eq!(
+			Err(Error::PersistenceFailed),
+			data_store.update_or_insert(new_update, new_object).await
+		);
+		assert!(data_store.get(&new_id).is_none());
+	}
+
 	#[tokio::test]
 	async fn insert_or_update_does_not_mutate_memory_if_persist_fails() {
 		let existing_id = TestObjectId { id: [42u8; 4] };

src/payment/pending_payment_store.rs

@@ -242,4 +242,78 @@ mod tests {
 			"current txid must not remain in its own conflict list"
 		);
 	}
+
+	#[test]
+	fn funding_classification_pending_update_preserves_mirrored_confirmation() {
+		use bitcoin::BlockHash;
+
+		use crate::payment::store::PaymentDetailsUpdate;
+
+		let txid = test_txid(7);
+		let payment_id = PaymentId(txid.to_byte_array());
+
+		// A pending entry wallet sync has already mirrored a confirmation into (via
+		// `apply_funding_status_update`) while classification was still writing its two stores.
+		let confirmed_details = PaymentDetails::new(
+			payment_id,
+			PaymentKind::Onchain {
+				txid,
+				status: ConfirmationStatus::Confirmed {
+					block_hash: BlockHash::from_byte_array([8u8; 32]),
+					height: 100,
+					timestamp: 1,
+				},
+				tx_type: None,
+			},
+			Some(2_000_000),
+			Some(999),
+			PaymentDirection::Outbound,
+			PaymentStatus::Pending,
+		);
+		let mirrored = PendingPaymentDetails::new(confirmed_details, Vec::new(), Vec::new());
+
+		// A fresh classification is always Unconfirmed and carries the candidate history; its
+		// figures are the active candidate's.
+		let fresh = pending_onchain_payment(payment_id, txid);
+		let candidates = vec![FundingTxCandidate {
+			txid,
+			amount_msat: fresh.amount_msat,
+			fee_paid_msat: fresh.fee_paid_msat,
+		}];
+
+		// The old fresh-insert path merged the full fresh record, downgrading the mirrored
+		// confirmation.
+		let mut downgraded = mirrored.clone();
+		let full_update =
+			PendingPaymentDetails::new(fresh.clone(), Vec::new(), candidates.clone()).to_update();
+		assert!(downgraded.update(full_update));
+		assert!(
+			matches!(
+				downgraded.details.kind,
+				PaymentKind::Onchain { status: ConfirmationStatus::Unconfirmed, .. }
+			),
+			"a full merge of a fresh classification downgrades a mirrored confirmation",
+		);
+
+		// The narrow classification update merges the figures and candidates while preserving the
+		// confirmation state wallet sync owns.
+		let mut merged = mirrored.clone();
+		let narrow_update = PendingPaymentDetailsUpdate {
+			id: payment_id,
+			payment_update: Some(PaymentDetailsUpdate::funding_reclassification(fresh)),
+			conflicting_txids: None,
+			candidates: candidates.clone(),
+		};
+		assert!(merged.update(narrow_update));
+		assert!(
+			matches!(
+				merged.details.kind,
+				PaymentKind::Onchain { status: ConfirmationStatus::Confirmed { .. }, .. }
+			),
+			"a narrow classification update must not downgrade a mirrored confirmation",
+		);
+		assert_eq!(merged.candidates, candidates);
+		assert_eq!(merged.details.amount_msat, Some(1_000));
+		assert_eq!(merged.details.fee_paid_msat, Some(100));
+	}
 }

src/payment/store.rs

@@ -1137,6 +1137,111 @@ mod tests {
 		assert_eq!(merged.fee_paid_msat, Some(500));
 	}
 
+	#[tokio::test]
+	async fn funding_classification_update_or_insert_preserves_advanced_record() {
+		use bitcoin::hashes::Hash;
+		use lightning::util::test_utils::TestLogger;
+		use std::str::FromStr;
+		use std::sync::Arc;
+
+		use crate::data_store::{DataStore, DataStoreUpdateOrInsertResult};
+		use crate::io::test_utils::InMemoryStore;
+		use crate::types::{DynStore, DynStoreWrapper};
+
+		let txid = Txid::from_byte_array([7u8; 32]);
+		let id = PaymentId(txid.to_byte_array());
+		let tx_type = Some(TransactionType::InteractiveFunding {
+			channels: vec![Channel {
+				counterparty_node_id: PublicKey::from_str(
+					"0279be667ef9dcbbac55a06295ce870b07029bfcdb2dce28d959f2815b16f81798",
+				)
+				.unwrap(),
+				channel_id: ChannelId([3u8; 32]),
+			}],
+		});
+		// A funding payment wallet sync has already advanced to Succeeded/Confirmed.
+		let advanced = PaymentDetails::new(
+			id,
+			PaymentKind::Onchain {
+				txid,
+				status: ConfirmationStatus::Confirmed {
+					block_hash: BlockHash::from_byte_array([8u8; 32]),
+					height: 100,
+					timestamp: 1,
+				},
+				tx_type: tx_type.clone(),
+			},
+			Some(2_000_000),
+			Some(999),
+			PaymentDirection::Outbound,
+			PaymentStatus::Succeeded,
+		);
+		// A fresh funding classification for the same payment is always Pending/Unconfirmed.
+		let fresh = PaymentDetails::new(
+			id,
+			PaymentKind::Onchain { txid, status: ConfirmationStatus::Unconfirmed, tx_type },
+			Some(1_000_000),
+			Some(500),
+			PaymentDirection::Outbound,
+			PaymentStatus::Pending,
+		);
+
+		let new_store = |seed: Vec<PaymentDetails>| {
+			let store: Arc<DynStore> = Arc::new(DynStoreWrapper(InMemoryStore::new()));
+			let logger = Arc::new(TestLogger::new());
+			DataStore::<PaymentDetails, Arc<TestLogger>>::new(
+				seed,
+				"payment_test_primary".to_string(),
+				"payment_test_secondary".to_string(),
+				store,
+				logger,
+			)
+		};
+
+		// The pre-fix fresh-insert path — a full `insert_or_update` merge landing after a racing
+		// wallet sync already advanced the record — downgrades it.
+		let store = new_store(vec![advanced.clone()]);
+		store.insert_or_update(fresh.clone()).await.unwrap();
+		let downgraded = store.get(&id).unwrap();
+		assert_eq!(
+			downgraded.status,
+			PaymentStatus::Pending,
+			"a full merge of a fresh classification downgrades an advanced record",
+		);
+
+		// `update_or_insert` applies only the narrow reclassification when a record exists — no
+		// matter when it appeared — preserving the confirmation state wallet sync owns while
+		// still merging the contribution-derived figures.
+		let store = new_store(vec![advanced.clone()]);
+		let update = PaymentDetailsUpdate::funding_reclassification(fresh.clone());
+		assert_eq!(
+			Ok(DataStoreUpdateOrInsertResult::Updated),
+			store.update_or_insert(update, fresh.clone()).await
+		);
+		let merged = store.get(&id).unwrap();
+		assert_eq!(merged.status, PaymentStatus::Succeeded);
+		assert!(matches!(
+			merged.kind,
+			PaymentKind::Onchain { status: ConfirmationStatus::Confirmed { .. }, .. }
+		));
+		assert_eq!(merged.amount_msat, Some(1_000_000));
+		assert_eq!(merged.fee_paid_msat, Some(500));
+
+		// And it inserts the fresh details when no record exists yet.
+		let store = new_store(Vec::new());
+		let update = PaymentDetailsUpdate::funding_reclassification(fresh.clone());
+		assert_eq!(
+			Ok(DataStoreUpdateOrInsertResult::Inserted),
+			store.update_or_insert(update, fresh).await
+		);
+		let inserted = store.get(&id).unwrap();
+		assert_eq!(inserted.status, PaymentStatus::Pending);
+		assert!(matches!(
+			inserted.kind,
+			PaymentKind::Onchain { status: ConfirmationStatus::Unconfirmed, .. }
+		));
+	}
+
 	#[derive(Clone, Debug, PartialEq, Eq)]
 	struct LegacyBolt11JitKind {
 		hash: PaymentHash,