Add PaginatedKVStore support to VssStore

Author: benthecarman

src/io/vss_store.rs

@@ -24,7 +24,7 @@ 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;
+use lightning::util::persist::{KVStore, PageToken, PaginatedKVStore, PaginatedListResponse};
 use lightning::util::ser::{Readable, Writeable};
 use prost::Message;
 use vss_client::client::VssClient;
@@ -293,6 +293,32 @@ impl KVStore for VssStore {
 	}
 }
 
+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);
+		let runtime = self.internal_runtime();
+		async move {
+			let task = runtime.spawn(async move {
+				inner
+					.list_paginated_internal(
+						&inner.async_client,
+						primary_namespace,
+						secondary_namespace,
+						page_token,
+					)
+					.await
+			});
+			task.await.map_err(|e| {
+				io::Error::new(io::ErrorKind::Other, format!("VSS runtime task failed: {}", e))
+			})?
+		}
+	}
+}
+
 impl Drop for VssStore {
 	fn drop(&mut self) {
 		if let Some(runtime) = self.internal_runtime.take() {
@@ -416,7 +442,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(
@@ -548,13 +576,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,
@@ -963,4 +1020,85 @@ mod tests {
 		do_read_write_remove_list_persist(&vss_store).await;
 		drop(vss_store)
 	}
+
+	#[tokio::test]
+	async 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];
+			KVStore::write(&store, ns, sub, &key, data).await.unwrap();
+		}
+
+		let mut all_keys = Vec::new();
+		let mut page_token = None;
+
+		loop {
+			let response =
+				PaginatedKVStore::list_paginated(&store, ns, sub, page_token).await.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);
+	}
+
+	#[tokio::test]
+	async fn vss_paginated_empty_namespace() {
+		let store = build_vss_store();
+		let response =
+			PaginatedKVStore::list_paginated(&store, "nonexistent", "ns", None).await.unwrap();
+		assert!(response.keys.is_empty());
+		assert!(response.next_page_token.is_none());
+	}
+
+	#[tokio::test]
+	async fn vss_paginated_removal() {
+		let store = build_vss_store();
+		let ns = "test_paginated";
+		let sub = "removal";
+
+		KVStore::write(&store, ns, sub, "a", vec![1u8; 8]).await.unwrap();
+		KVStore::write(&store, ns, sub, "b", vec![2u8; 8]).await.unwrap();
+		KVStore::write(&store, ns, sub, "c", vec![3u8; 8]).await.unwrap();
+
+		KVStore::remove(&store, ns, sub, "b", false).await.unwrap();
+
+		let response = PaginatedKVStore::list_paginated(&store, ns, sub, None).await.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()));
+	}
+
+	#[tokio::test]
+	async fn vss_paginated_namespace_isolation() {
+		let store = build_vss_store();
+
+		KVStore::write(&store, "ns_a", "sub", "key_1", vec![1u8; 8]).await.unwrap();
+		KVStore::write(&store, "ns_a", "sub", "key_2", vec![2u8; 8]).await.unwrap();
+		KVStore::write(&store, "ns_b", "sub", "key_3", vec![3u8; 8]).await.unwrap();
+
+		let response = PaginatedKVStore::list_paginated(&store, "ns_a", "sub", None).await.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 = PaginatedKVStore::list_paginated(&store, "ns_b", "sub", None).await.unwrap();
+		assert_eq!(response.keys.len(), 1);
+		assert!(response.keys.contains(&"key_3".to_string()));
+	}
 }