[stable2606] Backport #12106

prdoc/pr_12106.prdoc

@@ -0,0 +1,15 @@
+title: 'sp-database: aggregate same-key refcount ops within a transaction'
+doc:
+- audience: Node Dev
+  description: |
+    Fix a kvdb refcount-leak in `sp_database`: when one `Transaction` contains multiple
+    `Store`/`Reference`/`Release` ops on the same key, each op reads a stale on-disk
+    counter and writes back to the same counter key — the underlying `DBTransaction`
+    keeps only the last `put`, collapsing N ops to a single ±1. `commit_impl` now
+    replays per-`(col, key)` refcount ops in transaction order against one on-disk
+    counter read, then emits one final write per unique key.
+crates:
+- name: sc-client-db
+  bump: patch
+- name: sp-database
+  bump: patch

substrate/client/db/src/lib.rs

@@ -6744,6 +6744,135 @@ pub(crate) mod tests {
 			assert_eq!(get_value(&factory, h).as_deref(), Some(bytes.as_slice()));
 		}
 
+		#[rstest]
+		#[case::kvdb_memdb(BackendKind::KvdbMemdb)]
+		#[case::paritydb(BackendKind::ParityDb)]
+		#[case::rocksdb(BackendKind::RocksDb)]
+		fn release_then_store_missing_same_commit_stores_value(#[case] kind: BackendKind) {
+			let factory = DbFactory::new(kind);
+			let h = hash(0xAB);
+			let bytes = b"ab-bytes".to_vec();
+			{
+				let db = factory.open();
+				let mut tx = DbTransaction::new();
+				tx.release(TEST_COL, h);
+				tx.store(TEST_COL, h, bytes.clone());
+				db.commit(tx).unwrap();
+			}
+			assert_eq!(get_value(&factory, h).as_deref(), Some(bytes.as_slice()));
+			commit_release(&factory, h);
+			assert!(get_value(&factory, h).is_none(), "single release balances the store");
+		}
+
+		#[rstest]
+		#[case::kvdb_memdb(BackendKind::KvdbMemdb)]
+		#[case::paritydb(BackendKind::ParityDb)]
+		#[case::rocksdb(BackendKind::RocksDb)]
+		fn reference_then_store_missing_same_commit_single_release_removes_value(
+			#[case] kind: BackendKind,
+		) {
+			let factory = DbFactory::new(kind);
+			let h = hash(0xAC);
+			let bytes = b"ac-bytes".to_vec();
+			{
+				let db = factory.open();
+				let mut tx = DbTransaction::new();
+				tx.reference(TEST_COL, h);
+				tx.store(TEST_COL, h, bytes.clone());
+				db.commit(tx).unwrap();
+			}
+			assert_eq!(get_value(&factory, h).as_deref(), Some(bytes.as_slice()));
+			commit_release(&factory, h);
+			assert!(get_value(&factory, h).is_none(), "missing-key reference is a no-op");
+		}
+
+		#[rstest]
+		#[case::kvdb_memdb(BackendKind::KvdbMemdb)]
+		#[case::paritydb(BackendKind::ParityDb)]
+		#[case::rocksdb(BackendKind::RocksDb)]
+		fn release_then_reference_at_one_same_commit_removes_value(#[case] kind: BackendKind) {
+			let factory = DbFactory::new(kind);
+			let h = hash(0xAD);
+			let bytes = b"ad-bytes".to_vec();
+			commit_store(&factory, h, bytes);
+			{
+				let db = factory.open();
+				let mut tx = DbTransaction::new();
+				tx.release(TEST_COL, h);
+				tx.reference(TEST_COL, h);
+				db.commit(tx).unwrap();
+			}
+			assert!(get_value(&factory, h).is_none(), "reference after removal is a no-op");
+		}
+
+		// Same-commit multi-op refcount tests: each Store/Reference/Release must compose.
+
+		#[rstest]
+		#[case::kvdb_memdb(BackendKind::KvdbMemdb)]
+		#[case::paritydb(BackendKind::ParityDb)]
+		#[case::rocksdb(BackendKind::RocksDb)]
+		fn store_then_two_releases_same_commit_removes_value(#[case] kind: BackendKind) {
+			let factory = DbFactory::new(kind);
+			let h = hash(0xA8);
+			let bytes = b"a8-bytes".to_vec();
+			commit_store(&factory, h, bytes);
+			commit_reference(&factory, h);
+			assert!(get_value(&factory, h).is_some(), "rc=2 after store + reference");
+			{
+				let db = factory.open();
+				let mut tx = DbTransaction::new();
+				tx.release(TEST_COL, h);
+				tx.release(TEST_COL, h);
+				db.commit(tx).unwrap();
+			}
+			assert!(get_value(&factory, h).is_none(), "rc 2 -> 0 after two releases");
+		}
+
+		#[rstest]
+		#[case::kvdb_memdb(BackendKind::KvdbMemdb)]
+		#[case::paritydb(BackendKind::ParityDb)]
+		#[case::rocksdb(BackendKind::RocksDb)]
+		fn store_then_two_references_same_commit_increments_twice(#[case] kind: BackendKind) {
+			let factory = DbFactory::new(kind);
+			let h = hash(0xA9);
+			let bytes = b"a9-bytes".to_vec();
+			commit_store(&factory, h, bytes.clone());
+			{
+				let db = factory.open();
+				let mut tx = DbTransaction::new();
+				tx.reference(TEST_COL, h);
+				tx.reference(TEST_COL, h);
+				db.commit(tx).unwrap();
+			}
+			commit_release(&factory, h);
+			commit_release(&factory, h);
+			assert_eq!(
+				get_value(&factory, h).as_deref(),
+				Some(bytes.as_slice()),
+				"rc 3 -> 1 after two releases, value still present",
+			);
+			commit_release(&factory, h);
+			assert!(get_value(&factory, h).is_none(), "rc=0 after final release");
+		}
+
+		#[rstest]
+		#[case::kvdb_memdb(BackendKind::KvdbMemdb)]
+		#[case::paritydb(BackendKind::ParityDb)]
+		#[case::rocksdb(BackendKind::RocksDb)]
+		fn store_then_release_same_commit_net_zero_removes_value(#[case] kind: BackendKind) {
+			let factory = DbFactory::new(kind);
+			let h = hash(0xAA);
+			let bytes = b"aa-bytes".to_vec();
+			{
+				let db = factory.open();
+				let mut tx = DbTransaction::new();
+				tx.store(TEST_COL, h, bytes);
+				tx.release(TEST_COL, h);
+				db.commit(tx).unwrap();
+			}
+			assert!(get_value(&factory, h).is_none(), "store + release nets to rc=0");
+		}
+
 		struct BackendFactory {
 			backend: Option<Backend<Block>>,
 			kind: BackendKind,

substrate/primitives/database/src/kvdb.rs

@@ -17,6 +17,9 @@
 
 /// A wrapper around `kvdb::Database` that implements `sp_database::Database` trait
 use ::kvdb::{DBTransaction, KeyValueDB};
+use std::collections::HashMap;
+#[cfg(debug_assertions)]
+use std::collections::HashSet;
 
 use crate::{error, Change, ColumnId, Database, Transaction};
 
@@ -43,10 +46,10 @@ fn read_counter(
 		Some(data) => {
 			let mut counter_data = [0; 4];
 			if data.len() != 4 {
-				return Err(error::DatabaseError(Box::new(std::io::Error::new(
-					std::io::ErrorKind::Other,
-					format!("Unexpected counter len {}", data.len()),
-				))));
+				return Err(error::DatabaseError(Box::new(std::io::Error::other(format!(
+					"Unexpected counter len {}",
+					data.len(),
+				)))));
 			}
 			counter_data.copy_from_slice(&data);
 			let counter = u32::from_le_bytes(counter_data);
@@ -56,46 +59,117 @@ fn read_counter(
 	})
 }
 
+enum RefCountedOp {
+	Store(Vec<u8>),
+	Reference,
+	Release,
+}
+
 /// Commit a transaction to a KeyValueDB.
+///
+/// Ref-counted ops on the same `(col, key)` are replayed in order against one on-disk counter
+/// read, then the final counter/value state is emitted. Without this, multiple
+/// `Store`/`Reference`/`Release` in one tx would each read the stale on-disk counter and write
+/// back to the same counter key — the underlying batch keeps only the last `put`, collapsing N
+/// ops into one.
+///
+/// `Set`/`Remove` are emitted in submission order; ref-counted ops are emitted afterwards.
+/// Debug builds assert that raw and ref-counted ops are not mixed on the same `(col, key)`.
 fn commit_impl<H: Clone + AsRef<[u8]>>(
 	db: &dyn KeyValueDB,
 	transaction: Transaction<H>,
 ) -> error::Result<()> {
 	let mut tx = DBTransaction::new();
+	let mut ref_counted: HashMap<(ColumnId, Vec<u8>), Vec<RefCountedOp>> = HashMap::new();
+	#[cfg(debug_assertions)]
+	let mut raw_keys: HashSet<(ColumnId, Vec<u8>)> = HashSet::new();
+
 	for change in transaction.0.into_iter() {
 		match change {
-			Change::Set(col, key, value) => tx.put_vec(col, &key, value),
-			Change::Remove(col, key) => tx.delete(col, &key),
-			Change::Store(col, key, value) => match read_counter(db, col, key.as_ref())? {
-				(counter_key, Some(mut counter)) => {
-					counter += 1;
-					tx.put(col, &counter_key, &counter.to_le_bytes());
-				},
-				(counter_key, None) => {
-					let d = 1u32.to_le_bytes();
-					tx.put(col, &counter_key, &d);
-					tx.put_vec(col, key.as_ref(), value);
-				},
+			Change::Set(col, key, value) => {
+				#[cfg(debug_assertions)]
+				raw_keys.insert((col, key.clone()));
+				tx.put_vec(col, &key, value);
+			},
+			Change::Remove(col, key) => {
+				#[cfg(debug_assertions)]
+				raw_keys.insert((col, key.clone()));
+				tx.delete(col, &key);
+			},
+			Change::Store(col, key, value) => {
+				ref_counted
+					.entry((col, key.as_ref().to_vec()))
+					.or_default()
+					.push(RefCountedOp::Store(value));
 			},
 			Change::Reference(col, key) => {
-				if let (counter_key, Some(mut counter)) = read_counter(db, col, key.as_ref())? {
-					counter += 1;
-					tx.put(col, &counter_key, &counter.to_le_bytes());
-				}
+				ref_counted
+					.entry((col, key.as_ref().to_vec()))
+					.or_default()
+					.push(RefCountedOp::Reference);
 			},
 			Change::Release(col, key) => {
-				if let (counter_key, Some(mut counter)) = read_counter(db, col, key.as_ref())? {
-					counter -= 1;
-					if counter == 0 {
-						tx.delete(col, &counter_key);
-						tx.delete(col, key.as_ref());
-					} else {
-						tx.put(col, &counter_key, &counter.to_le_bytes());
+				ref_counted
+					.entry((col, key.as_ref().to_vec()))
+					.or_default()
+					.push(RefCountedOp::Release);
+			},
+		}
+	}
+
+	#[cfg(debug_assertions)]
+	for raw_key in &raw_keys {
+		debug_assert!(
+			!ref_counted.contains_key(raw_key),
+			"mixed raw/ref-counted database ops on column {}, key {:02x?}",
+			raw_key.0,
+			raw_key.1,
+		);
+	}
+
+	for ((col, key), ops) in ref_counted {
+		let (counter_key, mut counter) = read_counter(db, col, &key)?;
+
+		let mut value_to_write = None;
+		for op in ops {
+			match op {
+				RefCountedOp::Store(value) => match counter {
+					Some(c) => counter = Some(c + 1),
+					None => {
+						counter = Some(1);
+						value_to_write = Some(value);
+					},
+				},
+				RefCountedOp::Reference => {
+					if let Some(c) = counter {
+						counter = Some(c + 1);
 					}
+				},
+				RefCountedOp::Release => match counter {
+					Some(1) => {
+						counter = None;
+						value_to_write = None;
+					},
+					Some(c) => counter = Some(c - 1),
+					None => {},
+				},
+			}
+		}
+
+		match counter {
+			Some(counter) => {
+				tx.put(col, &counter_key, &counter.to_le_bytes());
+				if let Some(value) = value_to_write {
+					tx.put_vec(col, &key, value);
 				}
 			},
+			None => {
+				tx.delete(col, &counter_key);
+				tx.delete(col, &key);
+			},
 		}
 	}
+
 	db.write(tx).map_err(|e| error::DatabaseError(Box::new(e)))
 }