Merge pull request lightningdevkit#3481 from tnull/2024-12-add-kvstore-migration-ext

Add `MigratableKVStore` trait

Author: TheBlueMatt

lightning-persister/src/fs_store.rs

@@ -1,7 +1,7 @@
 //! Objects related to [`FilesystemStore`] live here.
 use crate::utils::{check_namespace_key_validity, is_valid_kvstore_str};
 
-use lightning::util::persist::KVStore;
+use lightning::util::persist::{KVStore, MigratableKVStore};
 use lightning::util::string::PrintableString;
 
 use std::collections::HashMap;
@@ -316,96 +316,187 @@ impl KVStore for FilesystemStore {
 			let entry = entry?;
 			let p = entry.path();
 
-			if let Some(ext) = p.extension() {
-				#[cfg(target_os = "windows")]
-				{
-					// Clean up any trash files lying around.
-					if ext == "trash" {
-						fs::remove_file(p).ok();
-						continue;
-					}
-				}
-				if ext == "tmp" {
-					continue;
-				}
+			if !dir_entry_is_key(&p)? {
+				continue;
 			}
 
-			let metadata = p.metadata()?;
+			let key = get_key_from_dir_entry(&p, &prefixed_dest)?;
 
-			// We allow the presence of directories in the empty primary namespace and just skip them.
-			if metadata.is_dir() {
-				continue;
+			keys.push(key);
+		}
+
+		self.garbage_collect_locks();
+
+		Ok(keys)
+	}
+}
+
+fn dir_entry_is_key(p: &Path) -> Result<bool, lightning::io::Error> {
+	if let Some(ext) = p.extension() {
+		#[cfg(target_os = "windows")]
+		{
+			// Clean up any trash files lying around.
+			if ext == "trash" {
+				fs::remove_file(p).ok();
+				return Ok(false);
 			}
+		}
+		if ext == "tmp" {
+			return Ok(false);
+		}
+	}
 
-			// If we otherwise don't find a file at the given path something went wrong.
-			if !metadata.is_file() {
+	let metadata = p.metadata().map_err(|e| {
+		let msg = format!(
+			"Failed to list keys at path {}: {}",
+			PrintableString(p.to_str().unwrap_or_default()),
+			e
+		);
+		lightning::io::Error::new(lightning::io::ErrorKind::Other, msg)
+	})?;
+
+	// We allow the presence of directories in the empty primary namespace and just skip them.
+	if metadata.is_dir() {
+		return Ok(false);
+	}
+
+	// If we otherwise don't find a file at the given path something went wrong.
+	if !metadata.is_file() {
+		debug_assert!(
+			false,
+			"Failed to list keys at path {}: file couldn't be accessed.",
+			PrintableString(p.to_str().unwrap_or_default())
+		);
+		let msg = format!(
+			"Failed to list keys at path {}: file couldn't be accessed.",
+			PrintableString(p.to_str().unwrap_or_default())
+		);
+		return Err(lightning::io::Error::new(lightning::io::ErrorKind::Other, msg));
+	}
+
+	Ok(true)
+}
+
+fn get_key_from_dir_entry(p: &Path, base_path: &Path) -> Result<String, lightning::io::Error> {
+	match p.strip_prefix(&base_path) {
+		Ok(stripped_path) => {
+			if let Some(relative_path) = stripped_path.to_str() {
+				if is_valid_kvstore_str(relative_path) {
+					return Ok(relative_path.to_string());
+				} else {
+					debug_assert!(
+						false,
+						"Failed to list keys of path {}: file path is not valid key",
+						PrintableString(p.to_str().unwrap_or_default())
+					);
+					let msg = format!(
+						"Failed to list keys of path {}: file path is not valid key",
+						PrintableString(p.to_str().unwrap_or_default())
+					);
+					return Err(lightning::io::Error::new(lightning::io::ErrorKind::Other, msg));
+				}
+			} else {
 				debug_assert!(
 					false,
-					"Failed to list keys of {}/{}: file couldn't be accessed.",
-					PrintableString(primary_namespace),
-					PrintableString(secondary_namespace)
+					"Failed to list keys of path {}: file path is not valid UTF-8",
+					PrintableString(p.to_str().unwrap_or_default())
 				);
 				let msg = format!(
-					"Failed to list keys of {}/{}: file couldn't be accessed.",
-					PrintableString(primary_namespace),
-					PrintableString(secondary_namespace)
+					"Failed to list keys of path {}: file path is not valid UTF-8",
+					PrintableString(p.to_str().unwrap_or_default())
 				);
 				return Err(lightning::io::Error::new(lightning::io::ErrorKind::Other, msg));
 			}
+		},
+		Err(e) => {
+			debug_assert!(
+				false,
+				"Failed to list keys of path {}: {}",
+				PrintableString(p.to_str().unwrap_or_default()),
+				e
+			);
+			let msg = format!(
+				"Failed to list keys of path {}: {}",
+				PrintableString(p.to_str().unwrap_or_default()),
+				e
+			);
+			return Err(lightning::io::Error::new(lightning::io::ErrorKind::Other, msg));
+		},
+	}
+}
+
+impl MigratableKVStore for FilesystemStore {
+	fn list_all_keys(&self) -> Result<Vec<(String, String, String)>, lightning::io::Error> {
+		let prefixed_dest = &self.data_dir;
+		if !prefixed_dest.exists() {
+			return Ok(Vec::new());
+		}
+
+		let mut keys = Vec::new();
+
+		'primary_loop: for primary_entry in fs::read_dir(prefixed_dest)? {
+			let primary_path = primary_entry?.path();
+
+			if dir_entry_is_key(&primary_path)? {
+				let primary_namespace = String::new();
+				let secondary_namespace = String::new();
+				let key = get_key_from_dir_entry(&primary_path, prefixed_dest)?;
+				keys.push((primary_namespace, secondary_namespace, key));
+				continue 'primary_loop;
+			}
+
+			// The primary_entry is actually also a directory.
+			'secondary_loop: for secondary_entry in fs::read_dir(&primary_path)? {
+				let secondary_path = secondary_entry?.path();
+
+				if dir_entry_is_key(&secondary_path)? {
+					let primary_namespace = get_key_from_dir_entry(&primary_path, prefixed_dest)?;
+					let secondary_namespace = String::new();
+					let key = get_key_from_dir_entry(&secondary_path, &primary_path)?;
+					keys.push((primary_namespace, secondary_namespace, key));
+					continue 'secondary_loop;
+				}
 
-			match p.strip_prefix(&prefixed_dest) {
-				Ok(stripped_path) => {
-					if let Some(relative_path) = stripped_path.to_str() {
-						if is_valid_kvstore_str(relative_path) {
-							keys.push(relative_path.to_string())
-						}
+				// The secondary_entry is actually also a directory.
+				for tertiary_entry in fs::read_dir(&secondary_path)? {
+					let tertiary_entry = tertiary_entry?;
+					let tertiary_path = tertiary_entry.path();
+
+					if dir_entry_is_key(&tertiary_path)? {
+						let primary_namespace =
+							get_key_from_dir_entry(&primary_path, prefixed_dest)?;
+						let secondary_namespace =
+							get_key_from_dir_entry(&secondary_path, &primary_path)?;
+						let key = get_key_from_dir_entry(&tertiary_path, &secondary_path)?;
+						keys.push((primary_namespace, secondary_namespace, key));
 					} else {
 						debug_assert!(
 							false,
-							"Failed to list keys of {}/{}: file path is not valid UTF-8",
-							PrintableString(primary_namespace),
-							PrintableString(secondary_namespace)
+							"Failed to list keys of path {}: only two levels of namespaces are supported",
+							PrintableString(tertiary_path.to_str().unwrap_or_default())
 						);
 						let msg = format!(
-							"Failed to list keys of {}/{}: file path is not valid UTF-8",
-							PrintableString(primary_namespace),
-							PrintableString(secondary_namespace)
+							"Failed to list keys of path {}: only two levels of namespaces are supported",
+							PrintableString(tertiary_path.to_str().unwrap_or_default())
 						);
 						return Err(lightning::io::Error::new(
 							lightning::io::ErrorKind::Other,
 							msg,
 						));
 					}
-				},
-				Err(e) => {
-					debug_assert!(
-						false,
-						"Failed to list keys of {}/{}: {}",
-						PrintableString(primary_namespace),
-						PrintableString(secondary_namespace),
-						e
-					);
-					let msg = format!(
-						"Failed to list keys of {}/{}: {}",
-						PrintableString(primary_namespace),
-						PrintableString(secondary_namespace),
-						e
-					);
-					return Err(lightning::io::Error::new(lightning::io::ErrorKind::Other, msg));
-				},
+				}
 			}
 		}
-
-		self.garbage_collect_locks();
-
 		Ok(keys)
 	}
 }
 
 #[cfg(test)]
 mod tests {
 	use super::*;
-	use crate::test_utils::{do_read_write_remove_list_persist, do_test_store};
+	use crate::test_utils::{
+		do_read_write_remove_list_persist, do_test_data_migration, do_test_store,
+	};
 
 	use bitcoin::Txid;
 
@@ -438,6 +529,19 @@ mod tests {
 		do_read_write_remove_list_persist(&fs_store);
 	}
 
+	#[test]
+	fn test_data_migration() {
+		let mut source_temp_path = std::env::temp_dir();
+		source_temp_path.push("test_data_migration_source");
+		let mut source_store = FilesystemStore::new(source_temp_path);
+
+		let mut target_temp_path = std::env::temp_dir();
+		target_temp_path.push("test_data_migration_target");
+		let mut target_store = FilesystemStore::new(target_temp_path);
+
+		do_test_data_migration(&mut source_store, &mut target_store);
+	}
+
 	#[test]
 	fn test_if_monitors_is_not_dir() {
 		let store = FilesystemStore::new("test_monitors_is_not_dir".into());

lightning-persister/src/test_utils.rs

@@ -3,7 +3,10 @@ use lightning::ln::functional_test_utils::{
 	connect_block, create_announced_chan_between_nodes, create_chanmon_cfgs, create_dummy_block,
 	create_network, create_node_cfgs, create_node_chanmgrs, send_payment,
 };
-use lightning::util::persist::{read_channel_monitors, KVStore, KVSTORE_NAMESPACE_KEY_MAX_LEN};
+use lightning::util::persist::{
+	migrate_kv_store_data, read_channel_monitors, KVStore, MigratableKVStore,
+	KVSTORE_NAMESPACE_KEY_ALPHABET, KVSTORE_NAMESPACE_KEY_MAX_LEN,
+};
 use lightning::util::test_utils;
 use lightning::{check_added_monitors, check_closed_broadcast, check_closed_event};
 
@@ -59,6 +62,41 @@ pub(crate) fn do_read_write_remove_list_persist<K: KVStore + RefUnwindSafe>(kv_s
 	assert_eq!(listed_keys.len(), 0);
 }
 
+pub(crate) fn do_test_data_migration<S: MigratableKVStore, T: MigratableKVStore>(
+	source_store: &mut S, target_store: &mut T,
+) {
+	// We fill the source with some bogus keys.
+	let dummy_data = [42u8; 32];
+	let num_primary_namespaces = 2;
+	let num_secondary_namespaces = 2;
+	let num_keys = 3;
+	for i in 0..num_primary_namespaces {
+		let primary_namespace =
+			format!("testspace{}", KVSTORE_NAMESPACE_KEY_ALPHABET.chars().nth(i).unwrap());
+		for j in 0..num_secondary_namespaces {
+			let secondary_namespace =
+				format!("testsubspace{}", KVSTORE_NAMESPACE_KEY_ALPHABET.chars().nth(j).unwrap());
+			for k in 0..num_keys {
+				let key =
+					format!("testkey{}", KVSTORE_NAMESPACE_KEY_ALPHABET.chars().nth(k).unwrap());
+				source_store
+					.write(&primary_namespace, &secondary_namespace, &key, &dummy_data)
+					.unwrap();
+			}
+		}
+	}
+	let total_num_entries = num_primary_namespaces * num_secondary_namespaces * num_keys;
+	let all_keys = source_store.list_all_keys().unwrap();
+	assert_eq!(all_keys.len(), total_num_entries);
+
+	migrate_kv_store_data(source_store, target_store).unwrap();
+
+	assert_eq!(target_store.list_all_keys().unwrap().len(), total_num_entries);
+	for (p, s, k) in &all_keys {
+		assert_eq!(target_store.read(p, s, k).unwrap(), dummy_data);
+	}
+}
+
 // Integration-test the given KVStore implementation. Test relaying a few payments and check that
 // the persisted data is updated the appropriate number of times.
 pub(crate) fn do_test_store<K: KVStore>(store_0: &K, store_1: &K) {

lightning/src/util/persist.rs

@@ -163,6 +163,40 @@ pub trait KVStore {
 	) -> Result<Vec<String>, io::Error>;
 }
 
+/// Provides additional interface methods that are required for [`KVStore`]-to-[`KVStore`]
+/// data migration.
+pub trait MigratableKVStore: KVStore {
+	/// Returns *all* known keys as a list of `primary_namespace`, `secondary_namespace`, `key` tuples.
+	///
+	/// This is useful for migrating data from [`KVStore`] implementation to [`KVStore`]
+	/// implementation.
+	///
+	/// Must exhaustively return all entries known to the store to ensure no data is missed, but
+	/// may return the items in arbitrary order.
+	fn list_all_keys(&self) -> Result<Vec<(String, String, String)>, io::Error>;
+}
+
+/// Migrates all data from one store to another.
+///
+/// This operation assumes that `target_store` is empty, i.e., any data present under copied keys
+/// might get overriden. User must ensure `source_store` is not modified during operation,
+/// otherwise no consistency guarantees can be given.
+///
+/// Will abort and return an error if any IO operation fails. Note that in this case the
+/// `target_store` might get left in an intermediate state.
+pub fn migrate_kv_store_data<S: MigratableKVStore, T: MigratableKVStore>(
+	source_store: &mut S, target_store: &mut T,
+) -> Result<(), io::Error> {
+	let keys_to_migrate = source_store.list_all_keys()?;
+
+	for (primary_namespace, secondary_namespace, key) in &keys_to_migrate {
+		let data = source_store.read(primary_namespace, secondary_namespace, key)?;
+		target_store.write(primary_namespace, secondary_namespace, key, &data)?;
+	}
+
+	Ok(())
+}
+
 /// Trait that handles persisting a [`ChannelManager`], [`NetworkGraph`], and [`WriteableScore`] to disk.
 ///
 /// [`ChannelManager`]: crate::ln::channelmanager::ChannelManager