StateKeyRegistry (aptos-labs#12864)

msmouse · web-flow · commit b58300b811b5 · 2024-04-30T21:52:16.000Z
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -450,6 +450,7 @@ aptos-cargo-cli = { path = "devtools/aptos-cargo-cli" }
 # External crate dependencies.
 # Please do not add any test features here: they should be declared by the individual crate.
 aes-gcm = "0.10.3"
+ahash = "0.8.11"
 atty = "0.2.14"
 nalgebra = "0.32"
 float-cmp = "0.9.0"
@@ -562,6 +563,7 @@ futures = "0.3.29"
 futures-channel = "0.3.29"
 futures-core = "0.3.29"
 futures-util = "0.3.29"
+fxhash = "0.2.1"
 getrandom = "0.2.2"
 gcp-bigquery-client = "0.13.0"
 get_if_addrs = "0.5.3"
diff --git a/api/types/src/wrappers.rs b/api/types/src/wrappers.rs
@@ -119,12 +119,7 @@ pub struct StateKeyWrapper(pub StateKey);
 
 impl fmt::Display for StateKeyWrapper {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        let hex_string = hex::encode(
-            self.0
-                .encode()
-                .context("Failed to encode StateKey")
-                .map_err(|_| fmt::Error)?,
-        );
+        let hex_string = hex::encode(self.0.encoded());
         write!(f, "{}", hex_string)
     }
 }
diff --git a/aptos-move/aptos-vm-types/src/storage/io_pricing.rs b/aptos-move/aptos-vm-types/src/storage/io_pricing.rs
@@ -61,12 +61,7 @@ impl IoPricingV1 {
         let mut cost = self.write_data_per_op * NumArgs::new(1);
 
         if self.write_data_per_byte_in_key > 0.into() {
-            cost += self.write_data_per_byte_in_key
-                * NumBytes::new(
-                    key.encode()
-                        .expect("Should be able to serialize state key")
-                        .len() as u64,
-                );
+            cost += self.write_data_per_byte_in_key * NumBytes::new(key.encoded().len() as u64);
         }
 
         match op_size {
@@ -135,11 +130,7 @@ impl IoPricingV2 {
                 .checked_sub(self.free_write_bytes_quota)
                 .unwrap_or(NumBytes::zero())
         } else {
-            let key_size = NumBytes::new(
-                key.encode()
-                    .expect("Should be able to serialize state key")
-                    .len() as u64,
-            );
+            let key_size = NumBytes::new(key.encoded().len() as u64);
             key_size + value_size
         }
     }
diff --git a/aptos-move/e2e-move-tests/src/tests/access_path_test.rs b/aptos-move/e2e-move-tests/src/tests/access_path_test.rs
@@ -18,8 +18,6 @@ use move_binary_format::{
 };
 use move_core_types::{identifier::Identifier, vm_status::StatusCode};
 
-/// FIXME(aldenhu): fix or remove
-#[ignore]
 #[test]
 fn access_path_panic() {
     // github.com/aptos-labs/aptos-core/security/advisories/GHSA-rpw2-84hq-48jj
diff --git a/storage/aptosdb/src/schema/stale_state_value_index/mod.rs b/storage/aptosdb/src/schema/stale_state_value_index/mod.rs
@@ -43,7 +43,7 @@ impl KeyCodec<StaleStateValueIndexSchema> for StaleStateValueIndex {
         let mut encoded = vec![];
         encoded.write_u64::<BigEndian>(self.stale_since_version)?;
         encoded.write_u64::<BigEndian>(self.version)?;
-        encoded.write_all(&self.state_key.encode()?)?;
+        encoded.write_all(self.state_key.encoded())?;
 
         Ok(encoded)
     }
diff --git a/storage/aptosdb/src/schema/state_value/mod.rs b/storage/aptosdb/src/schema/state_value/mod.rs
@@ -42,7 +42,7 @@ define_schema!(
 impl KeyCodec<StateValueSchema> for Key {
     fn encode_key(&self) -> Result<Vec<u8>> {
         let mut encoded = vec![];
-        encoded.write_all(&self.0.encode()?)?;
+        encoded.write_all(self.0.encoded())?;
         encoded.write_u64::<BigEndian>(!self.1)?;
         Ok(encoded)
     }
diff --git a/storage/aptosdb/src/schema/state_value_index/mod.rs b/storage/aptosdb/src/schema/state_value_index/mod.rs
@@ -30,7 +30,7 @@ define_schema!(StateValueIndexSchema, Key, (), STATE_VALUE_INDEX_CF_NAME);
 impl KeyCodec<StateValueIndexSchema> for Key {
     fn encode_key(&self) -> Result<Vec<u8>> {
         let mut encoded = vec![];
-        encoded.write_all(&self.0.encode()?)?;
+        encoded.write_all(self.0.encoded())?;
         encoded.write_u64::<BigEndian>(!self.1)?;
         Ok(encoded)
     }
diff --git a/testsuite/generate-format/Cargo.toml b/testsuite/generate-format/Cargo.toml
@@ -20,7 +20,7 @@ aptos-consensus-types = { workspace = true }
 aptos-crypto = { workspace = true }
 aptos-crypto-derive = { workspace = true }
 aptos-network = { workspace = true }
-aptos-types = { workspace = true }
+aptos-types = { workspace = true, features = ["fuzzing"] }
 bcs = { workspace = true }
 clap = { workspace = true }
 move-core-types = { workspace = true, features = ["fuzzing"] }
diff --git a/types/Cargo.toml b/types/Cargo.toml
@@ -19,6 +19,7 @@ aptos-crypto = { workspace = true }
 aptos-crypto-derive = { workspace = true }
 aptos-dkg = { workspace = true }
 aptos-experimental-runtimes = { workspace = true }
+aptos-infallible = { workspace = true }
 ark-bn254 = { workspace = true }
 ark-ff = { workspace = true }
 ark-groth16 = { workspace = true }
@@ -27,9 +28,9 @@ arr_macro = { workspace = true }
 base64 = { workspace = true }
 bcs = { workspace = true }
 bytes = { workspace = true }
-criterion = { workspace = true }
-derivative = { workspace = true }
 fixed = { workspace = true }
+fxhash = { workspace = true }
+hashbrown = { workspace = true }
 hex = { workspace = true }
 itertools = { workspace = true }
 jsonwebtoken = { workspace = true }
@@ -60,11 +61,15 @@ strum_macros = { workspace = true }
 thiserror = { workspace = true }
 
 [dev-dependencies]
+ahash = { workspace = true }
 aptos-crypto = { workspace = true, features = ["fuzzing", "testing"] }
+aptos-proptest-helpers = { workspace = true }
 async-trait = { workspace = true }
 ciborium = { workspace = true }
 claims = { workspace = true }
 coset = { workspace = true }
+criterion = { workspace = true }
+derivative = { workspace = true }
 move-core-types = { workspace = true, features = ["fuzzing"] }
 p256 = { workspace = true }
 passkey-authenticator = { workspace = true }
@@ -84,3 +89,7 @@ fuzzing = ["proptest", "proptest-derive", "aptos-crypto/fuzzing", "move-core-typ
 [[bench]]
 name = "keyless"
 harness = false
+
+[[bench]]
+name = "state_key"
+harness = false
diff --git a/types/benches/state_key.rs b/types/benches/state_key.rs
@@ -0,0 +1,196 @@
+// Copyright (c) Aptos Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+#![allow(clippy::non_canonical_partial_ord_impl)]
+
+use aptos_crypto::HashValue;
+use aptos_proptest_helpers::ValueGenerator;
+use aptos_types::{
+    access_path::AccessPath,
+    account_config::AccountResource,
+    state_store::state_key::{inner::StateKeyInner, registry::StateKeyRegistry, StateKey},
+};
+use criterion::{criterion_group, criterion_main, Criterion};
+use derivative::Derivative;
+use fxhash::FxHasher;
+use move_core_types::{
+    account_address::AccountAddress, language_storage::StructTag, move_resource::MoveStructType,
+};
+use once_cell::sync::OnceCell;
+use proptest::prelude::*;
+use std::{
+    collections::{hash_map::DefaultHasher, HashSet},
+    hash::{Hash, Hasher},
+};
+
+#[derive(Clone, Derivative, Eq)]
+#[derivative(Hash, PartialEq, PartialOrd, Ord)]
+pub struct UnCached {
+    inner: StateKeyInner,
+    #[derivative(
+        Hash = "ignore",
+        Ord = "ignore",
+        PartialEq = "ignore",
+        PartialOrd = "ignore"
+    )]
+    hash: OnceCell<HashValue>,
+}
+
+fn hashing(c: &mut Criterion) {
+    let mut group = c.benchmark_group("hashing");
+    let address = AccountAddress::from_hex_literal("0xA550C18").unwrap();
+    let struct_tag = AccountResource::struct_tag();
+
+    group.bench_function("hash_address_and_name", |b| {
+        b.iter(|| StateKeyRegistry::hash_address_and_name(&address, struct_tag.name.as_bytes()))
+    });
+
+    group.bench_function("default_hasher_struct_tag", |b| {
+        b.iter(|| {
+            let mut hasher = DefaultHasher::default();
+            Hash::hash(&struct_tag, &mut hasher);
+            hasher.finish()
+        })
+    });
+
+    group.bench_function("fxhasher_struct_tag", |b| {
+        b.iter(|| {
+            let mut hasher = FxHasher::default();
+            Hash::hash(&struct_tag, &mut hasher);
+            hasher.finish()
+        })
+    });
+
+    group.bench_function("ahasher_struct_tag", |b| {
+        b.iter(|| {
+            let mut hasher = ahash::AHasher::default();
+            Hash::hash(&struct_tag, &mut hasher);
+            hasher.finish()
+        })
+    });
+}
+
+fn construct(c: &mut Criterion) {
+    let keys = ValueGenerator::new().generate(proptest::collection::hash_set(
+        (any::<AccountAddress>(), any::<StructTag>()),
+        1_000,
+    ));
+
+    let mut group = c.benchmark_group("construct");
+    group.throughput(criterion::Throughput::Elements(keys.len() as u64));
+
+    group.bench_function("construct_once_uncached", |b| {
+        b.iter(|| {
+            keys.iter()
+                .map(|(address, struct_tag)| UnCached {
+                    inner: StateKeyInner::AccessPath(
+                        AccessPath::resource_access_path(*address, struct_tag.clone()).unwrap(),
+                    ),
+                    hash: OnceCell::new(),
+                })
+                .collect::<Vec<_>>()
+        })
+    });
+
+    group.bench_function("construct_once_cached", |b| {
+        b.iter(|| {
+            keys.iter()
+                .map(|(address, struct_tag)| StateKey::resource(address, struct_tag).unwrap())
+                .collect::<Vec<_>>()
+        })
+    });
+
+    group.bench_function("construct_twice_uncached", |b| {
+        b.iter(|| {
+            keys.iter()
+                .chain(keys.iter())
+                .map(|(address, struct_tag)| UnCached {
+                    inner: StateKeyInner::AccessPath(
+                        AccessPath::resource_access_path(*address, struct_tag.clone()).unwrap(),
+                    ),
+                    hash: OnceCell::new(),
+                })
+                .collect::<Vec<_>>()
+        })
+    });
+
+    group.bench_function("construct_twice_cached", |b| {
+        b.iter(|| {
+            keys.iter()
+                .chain(keys.iter())
+                .map(|(address, struct_tag)| StateKey::resource(address, struct_tag).unwrap())
+                .collect::<Vec<_>>()
+        })
+    });
+
+    group.bench_function("construct_thrice_uncached", |b| {
+        b.iter(|| {
+            keys.iter()
+                .chain(keys.iter())
+                .chain(keys.iter())
+                .map(|(address, struct_tag)| UnCached {
+                    inner: StateKeyInner::AccessPath(
+                        AccessPath::resource_access_path(*address, struct_tag.clone()).unwrap(),
+                    ),
+                    hash: OnceCell::new(),
+                })
+                .collect::<Vec<_>>()
+        })
+    });
+
+    group.bench_function("construct_thrice_cached", |b| {
+        b.iter(|| {
+            keys.iter()
+                .chain(keys.iter())
+                .chain(keys.iter())
+                .map(|(address, struct_tag)| StateKey::resource(address, struct_tag).unwrap())
+                .collect::<Vec<_>>()
+        })
+    });
+}
+
+fn hashset(c: &mut Criterion) {
+    let keys = ValueGenerator::new().generate(proptest::collection::hash_set(
+        any::<StateKeyInner>(),
+        1_000,
+    ));
+
+    let mut group = c.benchmark_group("hashset");
+    group.throughput(criterion::Throughput::Elements(keys.len() as u64));
+
+    let uncached = keys
+        .iter()
+        .map(|inner| UnCached {
+            inner: inner.clone(),
+            hash: OnceCell::new(),
+        })
+        .collect::<hashbrown::HashSet<_>>();
+    let cached = keys
+        .iter()
+        .map(|inner| StateKey::from_deserialized(inner.clone()).unwrap())
+        .collect::<hashbrown::HashSet<_>>();
+
+    group.bench_function("hashset_uncached", |b| {
+        b.iter(|| uncached.iter().cloned().collect::<HashSet<_>>())
+    });
+
+    group.bench_function("hashset_cached", |b| {
+        b.iter(|| cached.iter().cloned().collect::<HashSet<_>>())
+    });
+
+    group.bench_function("hashbrown_uncached", |b| {
+        b.iter(|| uncached.iter().cloned().collect::<hashbrown::HashSet<_>>())
+    });
+
+    group.bench_function("hashbrown_cached", |b| {
+        b.iter(|| cached.iter().cloned().collect::<hashbrown::HashSet<_>>())
+    });
+}
+
+criterion_group!(
+    name = state_key;
+    config = Criterion::default();
+    targets = hashing, construct, hashset
+);
+
+criterion_main!(state_key);
diff --git a/types/src/state_store/state_key/inner.rs b/types/src/state_store/state_key/inner.rs
diff --git a/types/src/state_store/state_key/mod.rs b/types/src/state_store/state_key/mod.rs
diff --git a/types/src/state_store/state_key/prefix.rs b/types/src/state_store/state_key/prefix.rs
diff --git a/types/src/state_store/state_key/registry.rs b/types/src/state_store/state_key/registry.rs

Original file line number	Diff line number	Diff line change
`@@ -119,12 +119,7 @@ pub struct StateKeyWrapper(pub StateKey);`
`119`	`119`
`120`	`120`	`impl fmt::Display for StateKeyWrapper {`
`121`	`121`	`fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {`
`122`		`- let hex_string = hex::encode(`
`123`		`- self.0`
`124`		`- .encode()`
`125`		`- .context("Failed to encode StateKey")`
`126`		`- .map_err(\|_\| fmt::Error)?,`
`127`		`- );`
	`122`	`+ let hex_string = hex::encode(self.0.encoded());`
`128`	`123`	`write!(f, "{}", hex_string)`
`129`	`124`	`}`
`130`	`125`	`}`
Original file line number	Diff line number	Diff line change
`@@ -43,7 +43,7 @@ impl KeyCodec<StaleStateValueIndexSchema> for StaleStateValueIndex {`
`43`	`43`	`let mut encoded = vec![];`
`44`	`44`	`encoded.write_u64::<BigEndian>(self.stale_since_version)?;`
`45`	`45`	`encoded.write_u64::<BigEndian>(self.version)?;`
`46`		`- encoded.write_all(&self.state_key.encode()?)?;`
	`46`	`+ encoded.write_all(self.state_key.encoded())?;`
`47`	`47`
`48`	`48`	`Ok(encoded)`
`49`	`49`	`}`
Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@ define_schema!(`
`42`	`42`	`impl KeyCodec<StateValueSchema> for Key {`
`43`	`43`	`fn encode_key(&self) -> Result<Vec<u8>> {`
`44`	`44`	`let mut encoded = vec![];`
`45`		`- encoded.write_all(&self.0.encode()?)?;`
	`45`	`+ encoded.write_all(self.0.encoded())?;`
`46`	`46`	`encoded.write_u64::<BigEndian>(!self.1)?;`
`47`	`47`	`Ok(encoded)`
`48`	`48`	`}`
Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,7 @@ define_schema!(StateValueIndexSchema, Key, (), STATE_VALUE_INDEX_CF_NAME);`
`30`	`30`	`impl KeyCodec<StateValueIndexSchema> for Key {`
`31`	`31`	`fn encode_key(&self) -> Result<Vec<u8>> {`
`32`	`32`	`let mut encoded = vec![];`
`33`		`- encoded.write_all(&self.0.encode()?)?;`
	`33`	`+ encoded.write_all(self.0.encoded())?;`
`34`	`34`	`encoded.write_u64::<BigEndian>(!self.1)?;`
`35`	`35`	`Ok(encoded)`
`36`	`36`	`}`