diff --git a/crates/blockchain/state_transition/src/beacon/helpers/test_state.rs b/crates/blockchain/state_transition/src/beacon/helpers/test_state.rs index 373e12262..d939fa20d 100644 --- a/crates/blockchain/state_transition/src/beacon/helpers/test_state.rs +++ b/crates/blockchain/state_transition/src/beacon/helpers/test_state.rs @@ -652,3 +652,42 @@ fn gloas_state(count: usize) -> gloas::BeaconState { .expect("the vector is built at its exact length"), } } + +#[cfg(test)] +mod tests { + use super::*; + + /// [`BeaconState::ssz_append_capacity`] for every fork. The test lives here + /// rather than beside that method in `ethlambda-types` because this module + /// is the one place a state of every fork is built. + /// + /// A small registry is the hard case. libssz's container encoder writes + /// each fixed-size field past the end of the buffer before moving it into + /// place, so a field needs room beyond the encoding's length unless enough + /// variable-size data still follows it, and a small state has little. A + /// fork that adds a fixed-size field larger than any the capacity counts + /// fails here rather than quietly regrowing a buffer the size of the state. + #[test] + fn ssz_append_never_regrows_a_buffer_with_its_capacity_to_spare() { + // A lean genesis with no validators: its fixed-size fields come before + // all of its variable-size ones, which hold almost nothing here. + let lean = BeaconState::Lean(ethlambda_types::state::State::from_genesis(7, vec![])); + let beacon = ForkName::ALL.map(|fork| with_validators_at(fork, 3)); + for state in beacon.into_iter().chain([lean]) { + let fork = state.fork_name(); + let capacity = state.ssz_append_capacity(); + + let encoded = state.to_ssz(); + assert_eq!(encoded.capacity(), capacity, "{fork}: to_ssz regrew"); + let decoded = BeaconState::from_ssz(fork, &encoded).expect("the encoding decodes"); + assert_eq!(decoded, state, "{fork}"); + + // Appending after a prefix reserves past the prefix, not from the + // start of the buffer. + let mut framed = vec![0xAB]; + state.ssz_append(&mut framed); + assert_eq!(framed.capacity(), 1 + capacity, "{fork}: ssz_append regrew"); + assert_eq!(framed[1..], encoded, "{fork}"); + } + } +} diff --git a/crates/common/types/src/beacon/containers/mod.rs b/crates/common/types/src/beacon/containers/mod.rs index a7859de2e..fc47ce679 100644 --- a/crates/common/types/src/beacon/containers/mod.rs +++ b/crates/common/types/src/beacon/containers/mod.rs @@ -366,8 +366,74 @@ impl BeaconState { /// unlike the accessors that read a beacon field: this is the inverse of /// [`BeaconState::from_ssz`], which builds that variant, so refusing here /// would make decoding a state and re-encoding it panic. + /// + /// Goes through [`BeaconState::ssz_append`] rather than libssz's own + /// `to_ssz`, which reserves only the encoding's length and so holds a fulu + /// or gloas state twice while encoding it (see + /// [`BeaconState::ssz_append_capacity`]). pub fn to_ssz(&self) -> Vec { - dispatch_state_including_lean!(self, |state| state.to_ssz()) + let mut buf = Vec::new(); + self.ssz_append(&mut buf); + buf + } + + /// How much [`BeaconState::ssz_append`] reserves past the end of the + /// buffer it appends to, so the append never regrows it: the encoding's + /// own length plus the largest fixed-size field of this variant's state. + /// Public for tests that check that bound. + /// + /// The headroom is for libssz's container encoder, which writes each + /// fixed-size field past the end of the buffer and then moves it into + /// place. A fixed-size field written after every variable-size one + /// (fulu's `proposer_lookahead`, gloas's `ptc_window`) therefore needs + /// room beyond the encoding's length, and without it a buffer the size of + /// the whole state is regrown to twice that, copying it, for the last few + /// hundred bytes. + pub fn ssz_append_capacity(&self) -> usize { + dispatch_state_including_lean!(self, |state| state.encoded_len()) + + self.largest_fixed_field() + } + + /// The largest fixed-size field in this variant's state: how far past the + /// encoding's end libssz's encoder can write (see + /// [`BeaconState::ssz_append_capacity`]). + /// + /// The beacon forks share one list. Every one of them has `randao_mixes`, + /// the largest on mainnet, and the only field that beats it on any preset + /// is gloas's `ptc_window` on minimal, by about a kilobyte, so a list per + /// fork would save nothing worth keeping eight lists for. A lean state is + /// a different container with fields a few hundred bytes wide, and + /// reserving a beacon state's megabytes for it would often outweigh the + /// state itself. + fn largest_fixed_field(&self) -> usize { + use libssz::SszEncode; + let largest = |sizes: &[usize]| sizes.iter().copied().max().unwrap_or(0); + match self { + BeaconState::Lean(_) => largest(&[ + ::fixed_size(), + ::fixed_size(), + ::fixed_size(), + ]), + _ => largest(&[ + ::fixed_size(), + ::fixed_size(), + ::fixed_size(), + ::fixed_size(), + ]), + } + } + + /// Appends [`BeaconState::to_ssz`]'s output to `buf`, after whatever it + /// already holds. + /// + /// For a caller that frames the encoding (the storage layer's fork tag): + /// a mainnet state encodes to hundreds of megabytes, and the append holds + /// one copy of it, never two. It first reserves + /// [`BeaconState::ssz_append_capacity`] past `buf`'s length, so a caller + /// need not size `buf` itself. + pub fn ssz_append(&self, buf: &mut Vec) { + buf.reserve_exact(self.ssz_append_capacity()); + dispatch_state_including_lean!(self, |state| state.ssz_append(buf)) } /// The state's merkle root, which a block's `state_root` must equal. diff --git a/crates/storage/src/api/traits.rs b/crates/storage/src/api/traits.rs index 35747be3c..e7f04643d 100644 --- a/crates/storage/src/api/traits.rs +++ b/crates/storage/src/api/traits.rs @@ -97,6 +97,16 @@ pub trait StorageWriteBatch: Send { /// Put multiple key-value pairs into a table. fn put_batch(&mut self, table: Table, batch: Vec<(Vec, Vec)>) -> Result<(), Error>; + /// Put one key-value pair into a table, from borrowed bytes. + /// + /// For a value the caller keeps after the write, so it need not clone one + /// only to hand it over: a beacon snapshot is hundreds of megabytes, and + /// RocksDB copies the value into its batch either way. The default goes + /// through [`put_batch`](Self::put_batch). + fn put(&mut self, table: Table, key: &[u8], value: &[u8]) -> Result<(), Error> { + self.put_batch(table, vec![(key.to_vec(), value.to_vec())]) + } + /// Delete multiple keys from a table. fn delete_batch(&mut self, table: Table, keys: Vec>) -> Result<(), Error>; diff --git a/crates/storage/src/backend/rocksdb.rs b/crates/storage/src/backend/rocksdb.rs index 4e51943c6..d74ad860a 100644 --- a/crates/storage/src/backend/rocksdb.rs +++ b/crates/storage/src/backend/rocksdb.rs @@ -219,6 +219,16 @@ impl StorageWriteBatch for RocksDBWriteBatch { Ok(()) } + fn put(&mut self, table: Table, key: &[u8], value: &[u8]) -> Result<(), Error> { + let cf = self + .db + .cf_handle(cf_name(table)) + .ok_or_else(|| format!("Column family {} not found", cf_name(table)))?; + // The batch copies the bytes, so the caller's buffer is never cloned. + self.batch.put_cf(&cf, key, value); + Ok(()) + } + fn delete_batch(&mut self, table: Table, keys: Vec>) -> Result<(), Error> { let cf = self .db diff --git a/crates/storage/src/backend/tests.rs b/crates/storage/src/backend/tests.rs index 0f9cd962c..6367f19bc 100644 --- a/crates/storage/src/backend/tests.rs +++ b/crates/storage/src/backend/tests.rs @@ -14,6 +14,7 @@ use crate::api::{StorageBackend, Table}; /// Run the full test suite against a backend. pub fn run_backend_tests(backend: &dyn StorageBackend) { test_put_and_get(backend); + test_put_from_borrowed_bytes(backend); test_delete(backend); test_prefix_iterator(backend); test_nonexistent_key(backend); @@ -47,6 +48,23 @@ fn test_put_and_get(backend: &dyn StorageBackend) { } } +fn test_put_from_borrowed_bytes(backend: &dyn StorageBackend) { + let value = b"borrowed_value".to_vec(); + { + let mut batch = backend.begin_write().unwrap(); + batch + .put(Table::BlockHeaders, b"test_put_borrowed_key", &value) + .unwrap(); + batch.commit().unwrap(); + } + + let view = backend.begin_read().unwrap(); + let stored = view + .get(Table::BlockHeaders, b"test_put_borrowed_key") + .unwrap(); + assert_eq!(stored, Some(value), "the caller's buffer is still its own"); +} + fn test_contains(backend: &dyn StorageBackend) { { let mut batch = backend.begin_write().unwrap(); diff --git a/crates/storage/src/state_codec.rs b/crates/storage/src/state_codec.rs index 3e0255d89..8278f6326 100644 --- a/crates/storage/src/state_codec.rs +++ b/crates/storage/src/state_codec.rs @@ -24,9 +24,11 @@ use ethlambda_types::{ /// [`ForkName::Lean`]'s selector is not a variant index, so the byte must go /// back through [`ForkName::from_selector`] rather than being cast. pub(crate) fn encode_state_value(state: &BeaconState) -> Vec { - let mut bytes = Vec::new(); - bytes.push(state.fork_name().selector()); - bytes.extend_from_slice(&state.to_ssz()); + // Encoded straight into the tagged buffer: the state writer runs this on + // every import, and a separate `Vec` would hold two copies of a mainnet + // state at once. `ssz_append` reserves its own room past the tag. + let mut bytes = vec![state.fork_name().selector()]; + state.ssz_append(&mut bytes); bytes } @@ -66,6 +68,22 @@ mod tests { assert_eq!(decode_state_value(&bytes), state); } + #[test] + fn a_beacon_state_value_is_its_tag_then_its_ssz_in_one_buffer() { + let electra = crate::beacon_state_delta::tests::electra_state_with_validators(64); + let state = BeaconState::Electra(electra); + let bytes = encode_state_value(&state); + + assert_eq!(bytes[0], ForkName::Electra.selector()); + assert_eq!(bytes[1..], state.to_ssz()); + assert_eq!( + bytes.capacity(), + 1 + state.ssz_append_capacity(), + "the buffer was regrown" + ); + assert_eq!(decode_state_value(&bytes), state); + } + #[test] fn the_selector_is_not_a_dense_index() { // ForkName::Lean is 255 so that beacon forks after fulu keep taking the diff --git a/crates/storage/src/state_writer.rs b/crates/storage/src/state_writer.rs index 51e5f6b0e..e22719ae0 100644 --- a/crates/storage/src/state_writer.rs +++ b/crates/storage/src/state_writer.rs @@ -81,12 +81,16 @@ pub(crate) struct StateWrite { /// `Table::StateDiffs` value. `Some` for every lean state, and for a /// beacon state that is not an anchor. pub diff: Option>, - /// Beacon only: the target's [`encode_state_value`] bytes, for the - /// writer's snapshot memo. `Some` exactly when the write is a snapshot: - /// only a snapshot is ever a diff base, so a diff's own bytes are never - /// kept. `None` on the lean arm, which diffs in the field domain and never - /// reads a base's bytes. - pub encoded: Option>, + /// Beacon only: `snapshot` also becomes the writer's snapshot memo, the + /// base the next diffs are taken against. True exactly when a beacon state + /// is written as a snapshot: only a snapshot is ever a diff base, so a + /// diff's own bytes are never kept. False on the lean arm, which diffs in + /// the field domain and never reads a base's bytes. + /// + /// A flag rather than a second copy of the bytes: the writer puts the + /// snapshot into its batch from a borrow and then moves that same buffer + /// into the memo, so a mainnet state is never held twice for it. + pub memoize_snapshot: bool, } /// States handed to the writer but not yet committed to the backend. @@ -504,7 +508,7 @@ pub(crate) fn plan_lean_state_write(state: State, parent: &State) -> StateWrite StateWrite { snapshot, diff: Some(diff), - encoded: None, + memoize_snapshot: false, } } @@ -531,9 +535,9 @@ pub(crate) fn plan_beacon_state_write( let target = encode_state_value(state); match base { None => StateWrite { - snapshot: Some(target.clone()), + snapshot: Some(target), diff: None, - encoded: Some(target), + memoize_snapshot: true, }, Some((snapshot_root, base)) => { let delta = beacon_state_delta::encode(&target, base); @@ -552,7 +556,7 @@ pub(crate) fn plan_beacon_state_write( StateWrite { snapshot: None, diff: Some(framed), - encoded: None, + memoize_snapshot: false, } } } @@ -693,17 +697,29 @@ impl StateWriter { .put_batch(Table::StateDiffs, vec![(key.clone(), diff)]) .expect("put state diff"); } - if let Some(snapshot) = write.snapshot { - batch - .put_batch(Table::States, vec![(key, snapshot)]) - .expect("put state snapshot"); - } + // A snapshot the memo keeps is put from a borrow, so the same buffer + // can then move into the memo; any other value is handed over. + let memo = match write.snapshot { + Some(snapshot) if write.memoize_snapshot => { + batch + .put(Table::States, &key, &snapshot) + .expect("put state snapshot"); + Some(snapshot) + } + Some(snapshot) => { + batch + .put_batch(Table::States, vec![(key, snapshot)]) + .expect("put state snapshot"); + None + } + None => None, + }; batch.commit().expect("commit"); if matches!(&*request.state, BeaconState::Lean(_)) { return; } - match write.encoded { + match memo { // A snapshot is its own base, and the next diffs are taken against it. Some(encoded) => { self.snapshot_memo = Some((root, encoded)); @@ -962,7 +978,7 @@ mod tests { Some(expected.as_slice()), "the snapshot holds the anchored state's own bytes, not its parent's" ); - assert!(plan.encoded.is_none(), "the lean arm keeps no encoded memo"); + assert!(!plan.memoize_snapshot, "the lean arm keeps no encoded memo"); let mut inside = base_state(); inside.slot = interval; @@ -986,10 +1002,9 @@ mod tests { "no base means nothing to diff against" ); assert!(plan.diff.is_none()); - assert_eq!( - plan.encoded.as_deref(), - plan.snapshot.as_deref(), - "the memo is the very bytes that were snapshotted" + assert!( + plan.memoize_snapshot, + "the snapshotted bytes become the memo" ); } @@ -1014,7 +1029,7 @@ mod tests { "a base means a diff, not a snapshot" ); assert!( - plan.encoded.is_none(), + !plan.memoize_snapshot, "only a snapshot is a diff base, so a diff keeps no memo bytes" ); let framed = plan.diff.expect("a non-anchor records a diff");