Add PaginatedKVStore support to VssStore

Author: benthecarman

src/io/vss_store.rs

@@ -24,7 +24,9 @@ use bitcoin::Network;
 use lightning::impl_writeable_tlv_based_enum;
 use lightning::io::{self, Error, ErrorKind};
 use lightning::sign::{EntropySource as LdkEntropySource, RandomBytes};
-use lightning::util::persist::{KVStore, KVStoreSync};
+use lightning::util::persist::{
+	KVStore, KVStoreSync, PageToken, PaginatedKVStore, PaginatedKVStoreSync, PaginatedListResponse,
+};
 use lightning::util::ser::{Readable, Writeable};
 use prost::Message;
 use vss_client::client::VssClient;
@@ -377,6 +379,52 @@ impl KVStore for VssStore {
 	}
 }
 
+impl PaginatedKVStoreSync for VssStore {
+	fn list_paginated(
+		&self, primary_namespace: &str, secondary_namespace: &str, page_token: Option<PageToken>,
+	) -> io::Result<PaginatedListResponse> {
+		let internal_runtime = self.internal_runtime.as_ref().ok_or_else(|| {
+			debug_assert!(false, "Failed to access internal runtime");
+			let msg = format!("Failed to access internal runtime");
+			Error::new(ErrorKind::Other, msg)
+		})?;
+		let primary_namespace = primary_namespace.to_string();
+		let secondary_namespace = secondary_namespace.to_string();
+		let inner = Arc::clone(&self.inner);
+		let fut = async move {
+			inner
+				.list_paginated_internal(
+					&inner.blocking_client,
+					primary_namespace,
+					secondary_namespace,
+					page_token,
+				)
+				.await
+		};
+		tokio::task::block_in_place(move || internal_runtime.block_on(fut))
+	}
+}
+
+impl PaginatedKVStore for VssStore {
+	fn list_paginated(
+		&self, primary_namespace: &str, secondary_namespace: &str, page_token: Option<PageToken>,
+	) -> impl Future<Output = Result<PaginatedListResponse, io::Error>> + 'static + Send {
+		let primary_namespace = primary_namespace.to_string();
+		let secondary_namespace = secondary_namespace.to_string();
+		let inner = Arc::clone(&self.inner);
+		async move {
+			inner
+				.list_paginated_internal(
+					&inner.async_client,
+					primary_namespace,
+					secondary_namespace,
+					page_token,
+				)
+				.await
+		}
+	}
+}
+
 impl Drop for VssStore {
 	fn drop(&mut self) {
 		let internal_runtime = self.internal_runtime.take();
@@ -499,7 +547,9 @@ impl VssStoreInner {
 			keys.push(self.extract_key(&kv.key)?);
 		}
 
-		Ok((keys, response.next_page_token))
+		// VSS may return an empty string instead of None to signal the last page.
+		let next_page_token = response.next_page_token.filter(|t| !t.is_empty());
+		Ok((keys, next_page_token))
 	}
 
 	async fn read_internal(
@@ -629,13 +679,42 @@ impl VssStoreInner {
 				.await?;
 			keys.extend(page_keys);
 			match next_page_token {
-				Some(t) if !t.is_empty() => page_token = Some(t),
-				_ => break,
+				Some(t) => page_token = Some(t),
+				None => break,
 			}
 		}
 		Ok(keys)
 	}
 
+	async fn list_paginated_internal(
+		&self, client: &VssClient<CustomRetryPolicy>, primary_namespace: String,
+		secondary_namespace: String, page_token: Option<PageToken>,
+	) -> io::Result<PaginatedListResponse> {
+		check_namespace_key_validity(
+			&primary_namespace,
+			&secondary_namespace,
+			None,
+			"list_paginated",
+		)?;
+
+		const PAGE_SIZE: i32 = 50;
+
+		let vss_page_token = page_token.map(|t| t.to_string());
+		let (keys, next_page_token) = self
+			.list_keys(
+				client,
+				&primary_namespace,
+				&secondary_namespace,
+				vss_page_token,
+				Some(PAGE_SIZE),
+			)
+			.await?;
+
+		let next_page_token = next_page_token.map(PageToken::new);
+
+		Ok(PaginatedListResponse { keys, next_page_token })
+	}
+
 	async fn execute_locked_write<
 		F: Future<Output = Result<(), lightning::io::Error>>,
 		FN: FnOnce() -> F,
@@ -1042,4 +1121,85 @@ mod tests {
 		do_read_write_remove_list_persist(&vss_store);
 		drop(vss_store)
 	}
+
+	#[test]
+	fn vss_paginated_listing() {
+		let store = build_vss_store();
+		let ns = "test_paginated";
+		let sub = "listing";
+		let num_entries = 5;
+
+		for i in 0..num_entries {
+			let key = format!("key_{:04}", i);
+			let data = vec![i as u8; 32];
+			KVStoreSync::write(&store, ns, sub, &key, data).unwrap();
+		}
+
+		let mut all_keys = Vec::new();
+		let mut page_token = None;
+
+		loop {
+			let response =
+				PaginatedKVStoreSync::list_paginated(&store, ns, sub, page_token).unwrap();
+			all_keys.extend(response.keys);
+			match response.next_page_token {
+				Some(token) => page_token = Some(token),
+				_ => break,
+			}
+		}
+
+		assert_eq!(all_keys.len(), num_entries);
+
+		// Verify no duplicates
+		let mut unique = all_keys.clone();
+		unique.sort();
+		unique.dedup();
+		assert_eq!(unique.len(), num_entries);
+	}
+
+	#[test]
+	fn vss_paginated_empty_namespace() {
+		let store = build_vss_store();
+		let response =
+			PaginatedKVStoreSync::list_paginated(&store, "nonexistent", "ns", None).unwrap();
+		assert!(response.keys.is_empty());
+		assert!(response.next_page_token.is_none());
+	}
+
+	#[test]
+	fn vss_paginated_removal() {
+		let store = build_vss_store();
+		let ns = "test_paginated";
+		let sub = "removal";
+
+		KVStoreSync::write(&store, ns, sub, "a", vec![1u8; 8]).unwrap();
+		KVStoreSync::write(&store, ns, sub, "b", vec![2u8; 8]).unwrap();
+		KVStoreSync::write(&store, ns, sub, "c", vec![3u8; 8]).unwrap();
+
+		KVStoreSync::remove(&store, ns, sub, "b", false).unwrap();
+
+		let response = PaginatedKVStoreSync::list_paginated(&store, ns, sub, None).unwrap();
+		assert_eq!(response.keys.len(), 2);
+		assert!(response.keys.contains(&"a".to_string()));
+		assert!(!response.keys.contains(&"b".to_string()));
+		assert!(response.keys.contains(&"c".to_string()));
+	}
+
+	#[test]
+	fn vss_paginated_namespace_isolation() {
+		let store = build_vss_store();
+
+		KVStoreSync::write(&store, "ns_a", "sub", "key_1", vec![1u8; 8]).unwrap();
+		KVStoreSync::write(&store, "ns_a", "sub", "key_2", vec![2u8; 8]).unwrap();
+		KVStoreSync::write(&store, "ns_b", "sub", "key_3", vec![3u8; 8]).unwrap();
+
+		let response = PaginatedKVStoreSync::list_paginated(&store, "ns_a", "sub", None).unwrap();
+		assert_eq!(response.keys.len(), 2);
+		assert!(response.keys.contains(&"key_1".to_string()));
+		assert!(response.keys.contains(&"key_2".to_string()));
+
+		let response = PaginatedKVStoreSync::list_paginated(&store, "ns_b", "sub", None).unwrap();
+		assert_eq!(response.keys.len(), 1);
+		assert!(response.keys.contains(&"key_3".to_string()));
+	}
 }