Implement PaginatedKVStore trait for SqliteStore and VssStore

Add paginated key listing support using PageToken for cursor-based pagination.

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

Author: benthecarman

src/io/sqlite_store/mod.rs

@@ -14,7 +14,9 @@ use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
 
 use lightning::io;
-use lightning::util::persist::{KVStore, KVStoreSync};
+use lightning::util::persist::{
+	KVStore, KVStoreSync, PageToken, PaginatedKVStore, PaginatedKVStoreSync, PaginatedListResponse,
+};
 use lightning_types::string::PrintableString;
 use rusqlite::{named_params, Connection};
 
@@ -34,7 +36,10 @@ pub const DEFAULT_SQLITE_DB_FILE_NAME: &str = "ldk_data.sqlite";
 pub const DEFAULT_KV_TABLE_NAME: &str = "ldk_data";
 
 // The current SQLite `user_version`, which we can use if we'd ever need to do a schema migration.
-const SCHEMA_USER_VERSION: u16 = 2;
+const SCHEMA_USER_VERSION: u16 = 3;
+
+// The default page size for paginated list operations.
+const DEFAULT_PAGE_SIZE: u32 = 100;
 
 /// A [`KVStoreSync`] implementation that writes to and reads from an [SQLite] database.
 ///
@@ -222,6 +227,33 @@ impl KVStoreSync for SqliteStore {
 	}
 }
 
+impl PaginatedKVStoreSync for SqliteStore {
+	fn list_paginated(
+		&self, primary_namespace: &str, secondary_namespace: &str, page_token: Option<PageToken>,
+	) -> io::Result<PaginatedListResponse> {
+		self.inner.list_paginated_internal(primary_namespace, secondary_namespace, page_token)
+	}
+}
+
+impl PaginatedKVStore for SqliteStore {
+	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 fut = tokio::task::spawn_blocking(move || {
+			inner.list_paginated_internal(&primary_namespace, &secondary_namespace, page_token)
+		});
+		async move {
+			fut.await.unwrap_or_else(|e| {
+				let msg = format!("Failed to IO operation due join error: {}", e);
+				Err(io::Error::new(io::ErrorKind::Other, msg))
+			})
+		}
+	}
+}
+
 struct SqliteStoreInner {
 	connection: Arc<Mutex<Connection>>,
 	data_dir: PathBuf,
@@ -289,7 +321,9 @@ impl SqliteStoreInner {
 			primary_namespace TEXT NOT NULL,
 			secondary_namespace TEXT DEFAULT \"\" NOT NULL,
 			key TEXT NOT NULL CHECK (key <> ''),
-			value BLOB, PRIMARY KEY ( primary_namespace, secondary_namespace, key )
+			value BLOB,
+			created_at INTEGER DEFAULT (strftime('%s','now')),
+			PRIMARY KEY ( primary_namespace, secondary_namespace, key )
 			);",
 			kv_table_name
 		);
@@ -299,6 +333,17 @@ impl SqliteStoreInner {
 			io::Error::new(io::ErrorKind::Other, msg)
 		})?;
 
+		// Create index for efficient paginated queries ordered by creation time, then key
+		let index_sql = format!(
+			"CREATE INDEX IF NOT EXISTS {}_created_at_idx ON {} (primary_namespace, secondary_namespace, created_at DESC, key ASC);",
+			kv_table_name, kv_table_name
+		);
+
+		connection.execute(&index_sql, []).map_err(|e| {
+			let msg = format!("Failed to create index on {}: {}", kv_table_name, e);
+			io::Error::new(io::ErrorKind::Other, msg)
+		})?;
+
 		let connection = Arc::new(Mutex::new(connection));
 		let write_version_locks = Mutex::new(HashMap::new());
 		Ok(Self { connection, data_dir, kv_table_name, write_version_locks })
@@ -472,6 +517,119 @@ impl SqliteStoreInner {
 		Ok(keys)
 	}
 
+	fn list_paginated_internal(
+		&self, primary_namespace: &str, secondary_namespace: &str, page_token: Option<PageToken>,
+	) -> io::Result<PaginatedListResponse> {
+		check_namespace_key_validity(
+			primary_namespace,
+			secondary_namespace,
+			None,
+			"list_paginated",
+		)?;
+
+		let locked_conn = self.connection.lock().unwrap();
+
+		// Extract the last_key from the page token
+		let last_key = page_token.map(|t| t.0);
+
+		// If last_key is provided, we need to get the created_at value for that key to use as a cursor
+		let last_created_at: Option<i64> = if let Some(ref key) = last_key {
+			let sql = format!(
+				"SELECT created_at FROM {} WHERE primary_namespace=:primary_namespace AND secondary_namespace=:secondary_namespace AND key=:key",
+				self.kv_table_name
+			);
+			let mut stmt = locked_conn.prepare_cached(&sql).map_err(|e| {
+				let msg = format!("Failed to prepare statement: {}", e);
+				io::Error::new(io::ErrorKind::Other, msg)
+			})?;
+
+			stmt.query_row(
+				named_params! {
+					":primary_namespace": primary_namespace,
+					":secondary_namespace": secondary_namespace,
+					":key": key,
+				},
+				|row| row.get(0),
+			)
+			.ok()
+		} else {
+			None
+		};
+
+		// Query for keys, ordered by created_at DESC, key ASC (newest first, with key as tiebreaker)
+		// For pagination with composite sort, we need: (created_at < last) OR (created_at = last AND key > last_key)
+		//
+		// We fetch DEFAULT_PAGE_SIZE + 1 rows as an optimization to detect if more pages exist.
+		// If we get more than DEFAULT_PAGE_SIZE rows, we know there are additional results beyond
+		// this page. We then discard the extra row and use the last returned key as the next page
+		// token. This avoids needing a separate COUNT(*) query to determine pagination state.
+		let (sql, params): (String, Vec<(&str, &dyn rusqlite::ToSql)>) = if let (
+			Some(ref created_at),
+			Some(ref key),
+		) =
+			(&last_created_at, &last_key)
+		{
+			(
+					format!(
+						"SELECT key FROM {} WHERE primary_namespace=:primary_namespace AND secondary_namespace=:secondary_namespace \
+						AND (created_at < :last_created_at OR (created_at = :last_created_at AND key > :last_key)) \
+						ORDER BY created_at DESC, key ASC LIMIT :limit",
+						self.kv_table_name
+					),
+					vec![
+						(":primary_namespace", &primary_namespace as &dyn rusqlite::ToSql),
+						(":secondary_namespace", &secondary_namespace as &dyn rusqlite::ToSql),
+						(":last_created_at", created_at as &dyn rusqlite::ToSql),
+						(":last_key", key as &dyn rusqlite::ToSql),
+						(":limit", &(DEFAULT_PAGE_SIZE + 1) as &dyn rusqlite::ToSql),
+					],
+				)
+		} else {
+			(
+					format!(
+						"SELECT key FROM {} WHERE primary_namespace=:primary_namespace AND secondary_namespace=:secondary_namespace \
+						ORDER BY created_at DESC, key ASC LIMIT :limit",
+						self.kv_table_name
+					),
+					vec![
+						(":primary_namespace", &primary_namespace as &dyn rusqlite::ToSql),
+						(":secondary_namespace", &secondary_namespace as &dyn rusqlite::ToSql),
+						(":limit", &(DEFAULT_PAGE_SIZE + 1) as &dyn rusqlite::ToSql),
+					],
+				)
+		};
+
+		let mut stmt = locked_conn.prepare_cached(&sql).map_err(|e| {
+			let msg = format!("Failed to prepare statement: {}", e);
+			io::Error::new(io::ErrorKind::Other, msg)
+		})?;
+
+		let mut keys = Vec::with_capacity((DEFAULT_PAGE_SIZE + 1) as usize);
+
+		let rows_iter = stmt.query_map(params.as_slice(), |row| row.get(0)).map_err(|e| {
+			let msg = format!("Failed to retrieve queried rows: {}", e);
+			io::Error::new(io::ErrorKind::Other, msg)
+		})?;
+
+		for k in rows_iter {
+			keys.push(k.map_err(|e| {
+				let msg = format!("Failed to retrieve queried rows: {}", e);
+				io::Error::new(io::ErrorKind::Other, msg)
+			})?);
+		}
+
+		// Check if we have more pages by seeing if we got the extra "peek" row.
+		// If so, remove it and generate a page token from the last actual result.
+		let next_page_token = if keys.len() > DEFAULT_PAGE_SIZE as usize {
+			keys.pop(); // Remove the extra "peek" row used to detect more pages
+			keys.last().cloned().map(PageToken)
+		} else {
+			None
+		};
+
+		Ok(PaginatedListResponse { keys, next_page_token })
+	}
+
 	fn execute_locked_write<F: FnOnce() -> Result<(), lightning::io::Error>>(
 		&self, inner_lock_ref: Arc<Mutex<u64>>, locking_key: String, version: u64, callback: F,
 	) -> Result<(), lightning::io::Error> {

src/io/test_utils.rs

@@ -5,10 +5,11 @@
 // http://opensource.org/licenses/MIT>, at your option. You may not use this file except in
 // accordance with one or both of these licenses.
 
-use std::collections::{hash_map, HashMap};
+use std::collections::HashMap;
 use std::future::Future;
 use std::panic::RefUnwindSafe;
 use std::path::PathBuf;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Mutex;
 
 use lightning::events::ClosureReason;
@@ -18,7 +19,8 @@ use lightning::ln::functional_test_utils::{
 	create_node_chanmgrs, send_payment, TestChanMonCfg,
 };
 use lightning::util::persist::{
-	KVStore, KVStoreSync, MonitorUpdatingPersister, KVSTORE_NAMESPACE_KEY_MAX_LEN,
+	KVStore, KVStoreSync, MonitorUpdatingPersister, PageToken, PaginatedKVStore,
+	PaginatedKVStoreSync, PaginatedListResponse, KVSTORE_NAMESPACE_KEY_MAX_LEN,
 };
 use lightning::util::test_utils;
 use lightning::{check_closed_broadcast, io};
@@ -36,14 +38,23 @@ type TestMonitorUpdatePersister<'a, K> = MonitorUpdatingPersister<
 
 const EXPECTED_UPDATES_PER_PAYMENT: u64 = 5;
 
+/// Entry in the in-memory store, tracking both data and creation order.
+#[derive(Debug, Clone)]
+struct StoreEntry {
+	data: Vec<u8>,
+	created_at: u64,
+}
+
 pub struct InMemoryStore {
-	persisted_bytes: Mutex<HashMap<String, HashMap<String, Vec<u8>>>>,
+	persisted_bytes: Mutex<HashMap<String, HashMap<String, StoreEntry>>>,
+	counter: AtomicU64,
 }
 
 impl InMemoryStore {
 	pub fn new() -> Self {
 		let persisted_bytes = Mutex::new(HashMap::new());
-		Self { persisted_bytes }
+		let counter = AtomicU64::new(0);
+		Self { persisted_bytes, counter }
 	}
 
 	fn read_internal(
@@ -53,9 +64,8 @@ impl InMemoryStore {
 		let prefixed = format!("{primary_namespace}/{secondary_namespace}");
 
 		if let Some(outer_ref) = persisted_lock.get(&prefixed) {
-			if let Some(inner_ref) = outer_ref.get(key) {
-				let bytes = inner_ref.clone();
-				Ok(bytes)
+			if let Some(entry) = outer_ref.get(key) {
+				Ok(entry.data.clone())
 			} else {
 				Err(io::Error::new(io::ErrorKind::NotFound, "Key not found"))
 			}
@@ -71,7 +81,15 @@ impl InMemoryStore {
 
 		let prefixed = format!("{primary_namespace}/{secondary_namespace}");
 		let outer_e = persisted_lock.entry(prefixed).or_insert(HashMap::new());
-		outer_e.insert(key.to_string(), buf);
+
+		// Only increment counter if this is a new key
+		let created_at = if let Some(existing) = outer_e.get(key) {
+			existing.created_at
+		} else {
+			self.counter.fetch_add(1, Ordering::SeqCst)
+		};
+
+		outer_e.insert(key.to_string(), StoreEntry { data: buf, created_at });
 		Ok(())
 	}
 
@@ -91,14 +109,54 @@ impl InMemoryStore {
 	fn list_internal(
 		&self, primary_namespace: &str, secondary_namespace: &str,
 	) -> io::Result<Vec<String>> {
-		let mut persisted_lock = self.persisted_bytes.lock().unwrap();
+		let persisted_lock = self.persisted_bytes.lock().unwrap();
 
 		let prefixed = format!("{primary_namespace}/{secondary_namespace}");
-		match persisted_lock.entry(prefixed) {
-			hash_map::Entry::Occupied(e) => Ok(e.get().keys().cloned().collect()),
-			hash_map::Entry::Vacant(_) => Ok(Vec::new()),
+		match persisted_lock.get(&prefixed) {
+			Some(entries) => Ok(entries.keys().cloned().collect()),
+			None => Ok(Vec::new()),
 		}
 	}
+
+	fn list_paginated_internal(
+		&self, primary_namespace: &str, secondary_namespace: &str, page_token: Option<PageToken>,
+	) -> io::Result<PaginatedListResponse> {
+		let persisted_lock = self.persisted_bytes.lock().unwrap();
+		let prefixed = format!("{primary_namespace}/{secondary_namespace}");
+
+		const PAGE_SIZE: usize = 100;
+
+		let entries = match persisted_lock.get(&prefixed) {
+			Some(e) => e,
+			None => return Ok(PaginatedListResponse { keys: Vec::new(), next_page_token: None }),
+		};
+
+		// Extract the last_key from the page token
+		let last_key = page_token.map(|t| t.0);
+
+		// Collect and sort by created_at descending (newest first)
+		let mut sorted_entries: Vec<(&String, &StoreEntry)> = entries.iter().collect();
+		sorted_entries
+			.sort_by(|a, b| b.1.created_at.cmp(&a.1.created_at).then_with(|| a.0.cmp(&b.0)));
+
+		// Find the starting index based on last_key
+		let start_idx = if let Some(ref key) = last_key {
+			sorted_entries.iter().position(|(k, _)| *k == key).map(|i| i + 1).unwrap_or(0)
+		} else {
+			0
+		};
+
+		// Get the next page of keys
+		let end_idx = std::cmp::min(start_idx + PAGE_SIZE, sorted_entries.len());
+		let keys: Vec<String> =
+			sorted_entries[start_idx..end_idx].iter().map(|(k, _)| (*k).clone()).collect();
+
+		// Determine if there are more pages
+		let next_page_token =
+			if end_idx < sorted_entries.len() { keys.last().cloned().map(PageToken) } else { None };
+
+		Ok(PaginatedListResponse { keys, next_page_token })
+	}
 }
 
 impl KVStore for InMemoryStore {
@@ -152,6 +210,28 @@ impl KVStoreSync for InMemoryStore {
 	}
 }
 
+impl PaginatedKVStoreSync for InMemoryStore {
+	fn list_paginated(
+		&self, primary_namespace: &str, secondary_namespace: &str, page_token: Option<PageToken>,
+	) -> io::Result<PaginatedListResponse> {
+		self.list_paginated_internal(primary_namespace, secondary_namespace, page_token)
+	}
+}
+
+impl PaginatedKVStore for InMemoryStore {
+	fn list_paginated(
+		&self, primary_namespace: &str, secondary_namespace: &str, page_token: Option<PageToken>,
+	) -> impl Future<Output = Result<PaginatedListResponse, io::Error>> + 'static + Send {
+		let res = PaginatedKVStoreSync::list_paginated(
+			self,
+			primary_namespace,
+			secondary_namespace,
+			page_token,
+		);
+		async move { res }
+	}
+}
+
 unsafe impl Sync for InMemoryStore {}
 unsafe impl Send for InMemoryStore {}