Persistent `LargeSmtForest::entries()` scans past the requested lineage prefix

[kkovaacs]

Packages versions

miden-crypto: 0.29.0

Bug description

We're seeing pretty bad performance with miden-node 0.15.2 on the 0.15 testnet, where a GetAccount call typically takes ~400ms for the MIDEN faucet account. Tracing data has revealed that it's the reconstruction of storage map contents from the account state forest (built upon LargeSmtForest) is what's taking a long time.

The MIDEN faucet account has four storage maps in its storage: two are empty, and two have exactly one element. Still, iterating over the entries in the corresponding LargeSmtForest tree takes ~100ms each.

Turns out that the RocksDB iterator code in the LargeSmtForest persistent backend is to blame. When iterating over leaves it does the following:

• LeafKey serializes the lineage before the leaf index (

  miden-vm/crates/crypto/src/merkle/smt/large_forest/backend/persistent/keys.rs

  Lines 16 to 32 in 7c587ef

  	pub struct LeafKey {
  	/// The lineage (and hence tree) to which the leaf belongs.
  	pub lineage: LineageId,
  	/// The logical index of the leaf within its parent tree.
  	pub index: u64,
  	}
  	impl Serializable for LeafKey {
  	fn write_into<W: ByteWriter>(&self, target: &mut W) {
  	target.write(self.lineage);
  	target.write(self.index);
  	}
  	fn get_size_hint(&self) -> usize {
  	size_of::<LineageId>() + size_of::<u64>()
  	}

  ), so RocksDB orders leaves by lineage prefix.

• PersistentBackend::entries() calls prefix_iterator_cf with lineage_bytes (

  miden-vm/crates/crypto/src/merkle/smt/large_forest/backend/persistent/mod.rs

  Lines 296 to 311 in 7c587ef

  	fn entries(&self, lineage: LineageId) -> Result<impl Iterator<Item = Result<TreeEntry>>> {
  	if !self.lineages.contains_key(&lineage) {
  	return Err(BackendError::UnknownLineage(lineage));
  	}
  	let lineage_bytes = lineage.to_bytes();
  	// In order to improve iteration performance significantly, we iterate with a prefix. As
  	// leaves are keyed on `LeafKey`, which begins with the bytes of the lineage, we can use the
  	// lineage as our prefix. That means that the iterator should only yield values whose key
  	// begins with the prefix with a high likelihood.
  	let pfx_iterator = self.db.prefix_iterator_cf(self.cf(LEAVES_CF)?, lineage_bytes);
  	// Data ownership concerns mean we cannot use this iterator directly even if we could change
  	// its type, so we delegate to our custom entries iterator impl.
  	Ok(PersistentBackendEntriesIterator::new(lineage, pfx_iterator))

  ).

• rocksdb::prefix_iterator_cf merely enables prefix_same_as_start and seeks forward from the prefix (https://github.com/rust-rocksdb/rust-rocksdb/blob/bb7d2168eab1bc7849f23adbcb825e3aba1bd2f4/src/db.rs#L1489-L1502).

• The wrapper documents that prefix_same_as_start is effective only when the column family has a non-null prefix_extractor (https://github.com/rust-rocksdb/rust-rocksdb/blob/bb7d2168eab1bc7849f23adbcb825e3aba1bd2f4/src/db_options.rs#L4011-L4022).

• The leaves column-family configuration does not install a prefix extractor (

  miden-vm/crates/crypto/src/merkle/smt/large_forest/backend/persistent/mod.rs

  Lines 1439 to 1484 in 7c587ef

  	// Now we set up our basic options for all column families.
  	let mut cf_opts = db::BlockBasedOptions::default();
  	cf_opts.set_block_cache(&cache);
  	cf_opts.set_bloom_filter(config.bloom_filter_bits, false);
  	cf_opts.set_whole_key_filtering(true); // Better for point lookups.
  	cf_opts.set_pin_l0_filter_and_index_blocks_in_cache(true); // Improves performance.
  	// From this, we can set up the configuration for each of our column families. We start with
  	// the one for metadata.
  	let metadata_cf_opts = Self::build_cf_opts(
  	config,
  	&cf_opts,
  	MAX_METADATA_CF_WRITE_BUFFER_SIZE_BYTES,
  	db::DBCompressionType::None,
  	);
  	// We can also create the configuration for our leaves column family.
  	let leaves_cf_opts = Self::build_cf_opts(
  	config,
  	&cf_opts,
  	MAX_LEAVES_CF_WRITE_BUFFER_SIZE_BYTES,
  	COMPRESSION_MODE,
  	);
  	// Finally we create them for each of our subtree CFs.
  	let subtree_cfs = SUBTREE_CFS.into_iter().map(|name| {
  	db::ColumnFamilyDescriptor::new(
  	name,
  	Self::build_cf_opts(
  	config,
  	&cf_opts,
  	MAX_SUBTREE_CF_WRITE_BUFFER_SIZE_BYTES,
  	COMPRESSION_MODE,
  	),
  	)
  	});
  	// With the column-specific configuration made, we can then simply create our database
  	// options
  	let mut columns = vec![
  	db::ColumnFamilyDescriptor::new(METADATA_CF, metadata_cf_opts),
  	db::ColumnFamilyDescriptor::new(LEAVES_CF, leaves_cf_opts),
  	];
  	columns.extend(subtree_cfs);
  	Ok(DB::open_cf_descriptors(&db_opts, config.path.clone(), columns)?)

  ).

• RocksDB internally disables prefix_same_as_start when prefix_extractor_ is null (https://github.com/facebook/rocksdb/blob/410c5623195ecbe4699b9b5a5f622c7325cec6fe/db/db_iter.cc#L78-L83).

• The custom iterator deserializes each key and uses continue when the lineage differs (

  miden-vm/crates/crypto/src/merkle/smt/large_forest/backend/persistent/iterator.rs

  Lines 81 to 106 in 7c587ef

  	let entry = self.iterator.next()?;
  	let (key_bytes, value_bytes) = match entry {
  	Ok((key_bytes, value_bytes)) => (key_bytes, value_bytes),
  	Err(e) => {
  	self.state = PersistentBackendEntriesIteratorState::Faulted;
  	return Some(Err(e.into()));
  	},
  	};
  	let key = match LeafKey::read_from_bytes(&key_bytes) {
  	Ok(key) => key,
  	Err(e) => {
  	self.state = PersistentBackendEntriesIteratorState::Faulted;
  	return Some(Err(e.into()));
  	},
  	};
  	// If the key isn't for the correct lineage (which can happen even with the
  	// bloom filter), we need to advance by returning to the loop.
  	if key.lineage != self.lineage {
  	continue;
  	}
  	// If the key is valid, we need to read out the leaf itself and then start
  	// iterating over that.
  	let leaf = match SmtLeaf::read_from_bytes(&value_bytes) {

  ), causing it to consume later keys until the underlying iterator reaches the column-family end.

• The snapshot-backed reader independently enables the same ineffective option without an upper bound (

  miden-vm/crates/crypto/src/merkle/smt/large_forest/backend/persistent/snapshot.rs

  Lines 205 to 218 in 7c587ef

  	fn entries(&self, lineage: LineageId) -> Result<impl Iterator<Item = Result<TreeEntry>>> {
  	if !self.inner.lineages.contains_key(&lineage) {
  	return Err(BackendError::UnknownLineage(lineage));
  	}
  	let lineage_bytes = lineage.to_bytes();
  	let cf = self.cf(LEAVES_CF)?;
  	let mut read_opts = db::ReadOptions::default();
  	read_opts.set_prefix_same_as_start(true);
  	let pfx_iterator = self.inner.snapshot.iterator_cf_opt(
  	cf,
  	read_opts,
  	db::IteratorMode::From(&lineage_bytes, db::Direction::Forward),
  	);
  	Ok(PersistentBackendEntriesIterator::new(lineage, pfx_iterator))

  ), so it has the same issue.

How can this be reproduced?

#[test]
fn entries_stops_at_end_of_lineage_prefix() -> Result<()> {
    let (_file, mut backend) = default_backend()?;

    // Choose the lowest possible lineage so that any key beginning with a nonzero byte sorts after
    // all of this lineage's leaves.
    let lineage = LineageId::new([0; 32]);
    let key = Word::from([1_u32, 2, 3, 4]);
    let value = Word::from([5_u32, 6, 7, 8]);
    backend.add_lineage(lineage, 1, SmtUpdateBatch::from([(key, value)].into_iter()))?;

    // Insert a malformed key outside the requested lineage's prefix. A correctly bounded entries
    // iterator must stop before inspecting it. An unbounded iterator will try to deserialize it as
    // a LeafKey and return an error.
    let leaves_cf = backend.cf(LEAVES_CF)?;
    backend.db.put_cf(leaves_cf, [1], [])?;

    let entries = backend.entries(lineage)?.collect::<Result<Vec<_>>>()?;
    assert_eq!(entries, vec![TreeEntry { key, value }]);

    Ok(())
}

Relevant log output