refactor: extract ConsignmentCache

bitlightlabs · Oct 9, 2024 · 22bf975 · 22bf975
1 parent e8293ee
commit 22bf975
Show file tree

Hide file tree

Showing 2 changed files with 97 additions and 28 deletions.
diff --git a/lightning/src/color_ext/database.rs b/lightning/src/color_ext/database.rs
@@ -1,6 +1,7 @@
 use core::fmt::{Display, Formatter};
 use std::{
 	collections::HashMap,
+	io::{self, Write},
 	sync::{Arc, Mutex},
 };
 
@@ -198,11 +199,88 @@ impl RgbInfoCache {
 	}
 }
 
+#[derive(Clone, Debug, Default)]
+pub struct ConsignmentBinaryData(Vec<u8>);
+
+impl Write for ConsignmentBinaryData {
+	fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+		self.0.extend_from_slice(buf);
+		Ok(buf.len())
+	}
+
+	fn flush(&mut self) -> io::Result<()> {
+		Ok(())
+	}
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct ConsignmentHandle(usize);
+
+#[derive(Clone, Debug, Default)]
+struct ConsignmentCache {
+    by_channel_id: HashMap<ChannelId, ConsignmentHandle>,
+    by_funding_txid: HashMap<Txid, ConsignmentHandle>,
+    data_store: HashMap<ConsignmentHandle, ConsignmentBinaryData>,
+    next_handle: usize,
+}
+
+impl ConsignmentCache {
+    fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn get_by_channel_id(&self, channel_id: &ChannelId) -> Option<ConsignmentHandle> {
+        self.by_channel_id.get(channel_id).copied()
+    }
+
+    pub fn get_by_funding_txid(&self, funding_txid: &Txid) -> Option<ConsignmentHandle> {
+        self.by_funding_txid.get(funding_txid).copied()
+    }
+
+    pub fn insert(&mut self, channel_id: &ChannelId, funding_txid: Txid, info: ConsignmentBinaryData)-> ConsignmentHandle {
+        let handle = ConsignmentHandle(self.next_handle);
+        self.next_handle += 1;
+        self.data_store.insert(handle, info);
+        self.by_channel_id.insert(channel_id.clone(), handle);
+        self.by_funding_txid.insert(funding_txid, handle);
+
+		handle
+    }
+
+    pub fn remove(&mut self, channel_id: &ChannelId, funding_txid: Txid) {
+        if let Some(handle) = self.by_channel_id.remove(channel_id) {
+            self.by_funding_txid.retain(|_, &mut v| v != handle);
+            self.data_store.remove(&handle);
+        }
+        self.by_funding_txid.remove(&funding_txid);
+    }
+
+    pub fn resolve(&self, handle: ConsignmentHandle) -> Option<&ConsignmentBinaryData> {
+        self.data_store.get(&handle)
+    }
+
+	pub fn rename_channel_id(&mut self, handle: ConsignmentHandle, old_channel_id: &ChannelId, new_channel_id: &ChannelId) {
+		if let Some(info) = self.data_store.get(&handle) {
+			let mut new_info = ConsignmentBinaryData(Vec::new());
+			new_info.write_all(info.0.as_slice()).unwrap();
+			new_info.0.extend_from_slice(new_channel_id.as_ref());
+			self.data_store.insert(handle, new_info);
+		}
+
+		if let Some(old_handle) = self.by_channel_id.remove(old_channel_id) {
+			if old_handle == handle {
+				self.by_channel_id.insert(new_channel_id.clone(), handle);
+			}
+		}
+	}
+}
+
 #[derive(Default)]
 pub struct ColorDatabaseImpl {
 	rgb_payment_cache: Arc<Mutex<RgbPaymentCache>>,
 	transfer_info: Arc<Mutex<TransferInfoCache>>,
 	rgb_info: Arc<Mutex<RgbInfoCache>>,
+	consignment_cache: Arc<Mutex<ConsignmentCache>>,
 }
 
 impl ColorDatabaseImpl {
@@ -222,8 +300,12 @@ impl ColorDatabaseImpl {
 		self.rgb_info.clone()
 	}
 
+	pub fn consignment(&self) -> Arc<Mutex<ConsignmentCache>> {
+		self.consignment_cache.clone()
+	}
+
 	pub fn rename_channel_id(&self, old_channel_id: &ChannelId, new_channel_id: &ChannelId) {
-		let mut rgb_info_key = RgbInfoKey::new(old_channel_id, false);
+		let rgb_info_key = RgbInfoKey::new(old_channel_id, false);
 		if let Some(info) = self.rgb_info().lock().unwrap().get_by_rgb_info_key(&rgb_info_key) {
 			let new_info = info.clone();
 			self.rgb_info()
@@ -233,7 +315,7 @@ impl ColorDatabaseImpl {
 			self.rgb_info().lock().unwrap().remove(&rgb_info_key);
 		}
 
-		let mut rgb_info_key_pending = RgbInfoKey::new(old_channel_id, true);
+		let rgb_info_key_pending = RgbInfoKey::new(old_channel_id, true);
 		if let Some(info) =
 			self.rgb_info().lock().unwrap().get_by_rgb_info_key(&rgb_info_key_pending)
 		{
@@ -242,6 +324,8 @@ impl ColorDatabaseImpl {
 			self.rgb_info().lock().unwrap().remove(&rgb_info_key_pending);
 		}
 
-		// todo: rename consignment
+		if let Some(consignment_handle) = self.consignment().lock().unwrap().get_by_channel_id(old_channel_id) {
+			self.consignment().lock().unwrap().rename_channel_id(consignment_handle, old_channel_id, new_channel_id);
+		}
 	}
 }
diff --git a/lightning/src/color_ext/mod.rs b/lightning/src/color_ext/mod.rs
@@ -12,8 +12,8 @@ use crate::sign::SignerProvider;
 use bitcoin::blockdata::transaction::Transaction;
 use bitcoin::psbt::{PartiallySignedTransaction, Psbt};
 use bitcoin::secp256k1::PublicKey;
-use bitcoin::TxOut;
-use database::{ColorDatabaseImpl, PaymentHashKey, ProxyIdKey, RgbInfoKey};
+use bitcoin::{TxOut, Txid};
+use database::{ColorDatabaseImpl, ConsignmentBinaryData, PaymentHashKey, ProxyIdKey, RgbInfoKey};
 use hex::DisplayHex;
 use rgb_lib::wallet::rust_only::AssetBeneficiariesMap;
 use rgb_lib::Fascia;
@@ -517,17 +517,10 @@ impl ColorSourceImpl {
 	}
 
 	/// Write RgbInfo file
-	pub fn write_rgb_channel_info(&self, key: &RgbInfoKey, rgb_info: &RgbInfo) {
+	pub fn save_rgb_channel_info(&self, key: &RgbInfoKey, rgb_info: &RgbInfo) {
 		self.database.rgb_info().lock().unwrap().insert(*key, rgb_info.clone());
 	}
 
-	fn _append_pending_extension(&self, path: &Path) -> PathBuf {
-		let mut new_path = path.to_path_buf();
-		new_path
-			.set_extension(format!("{}_pending", new_path.extension().unwrap().to_string_lossy()));
-		new_path
-	}
-
 	/// Rename RGB files from temporary to final channel ID
 	pub(crate) fn rename_rgb_files(
 		&self, channel_id: &ChannelId, temporary_channel_id: &ChannelId,
@@ -536,13 +529,6 @@ impl ColorSourceImpl {
 		let chan_id = channel_id;
 
 		self.database.rename_channel_id(temp_chan_id, chan_id);
-
-		let funding_consignment_tmp =
-			self.ldk_data_dir.join(format!("consignment_{}", temp_chan_id));
-		if funding_consignment_tmp.exists() {
-			let funding_consignment = self.ldk_data_dir.join(format!("consignment_{}", chan_id));
-			fs::rename(funding_consignment_tmp, funding_consignment).expect("rename ok");
-		}
 	}
 
 	/// Handle funding on the receiver side
@@ -583,19 +569,18 @@ impl ColorSourceImpl {
 			},
 		};
 
-		let consignment_path = self.ldk_data_dir.join(format!("consignment_{}", funding_txid));
-		consignment.save_file(consignment_path).expect("unable to write file");
-		let consignment_path =
-			self.ldk_data_dir.join(format!("consignment_{}", temporary_channel_id.0.as_hex()));
-		consignment.save_file(consignment_path).expect("unable to write file");
+		let funding_txid = Txid::from_str(&funding_txid).unwrap();
+		let mut consignment_data = ConsignmentBinaryData::default();
+		consignment.save(&mut consignment_data);
+		self.database.consignment().lock().unwrap().insert(temporary_channel_id, funding_txid, consignment_data);
 
 		let rgb_info = RgbInfo {
 			contract_id: consignment.contract_id(),
 			local_rgb_amount: 0,
 			remote_rgb_amount,
 		};
-		self.write_rgb_channel_info(&RgbInfoKey::new(&temporary_channel_id, true), &rgb_info);
-		self.write_rgb_channel_info(&RgbInfoKey::new(&temporary_channel_id, false), &rgb_info);
+		self.save_rgb_channel_info(&RgbInfoKey::new(&temporary_channel_id, true), &rgb_info);
+		self.save_rgb_channel_info(&RgbInfoKey::new(&temporary_channel_id, false), &rgb_info);
 
 		Ok(())
 	}
@@ -620,7 +605,7 @@ impl ColorSourceImpl {
 			rgb_info.remote_rgb_amount -= received;
 		}
 
-		self.write_rgb_channel_info(&key, &rgb_info)
+		self.save_rgb_channel_info(&key, &rgb_info)
 	}
 
 	/// Update pending RGB channel amount