diff --git a/crates/stages/types/src/id.rs b/crates/stages/types/src/id.rs index 86dd9ced5c7..78d7e0ec1b6 100644 --- a/crates/stages/types/src/id.rs +++ b/crates/stages/types/src/id.rs @@ -1,3 +1,7 @@ +use alloc::vec::Vec; +#[cfg(feature = "std")] +use std::{collections::HashMap, sync::OnceLock}; + /// Stage IDs for all known stages. /// /// For custom stages, use [`StageId::Other`] @@ -27,6 +31,12 @@ pub enum StageId { Other(&'static str), } +/// One-time-allocated stage ids encoded as raw Vecs, useful for database +/// clients to reference them for queries instead of encoding anew per query +/// (sad heap allocation required). +#[cfg(feature = "std")] +static ENCODED_STAGE_IDS: OnceLock>> = OnceLock::new(); + impl StageId { /// All supported Stages pub const ALL: [Self; 15] = [ @@ -98,6 +108,25 @@ impl StageId { pub const fn is_finish(&self) -> bool { matches!(self, Self::Finish) } + + /// Get a pre-encoded raw Vec, for example, to be used as the DB key for + /// `tables::StageCheckpoints` and `tables::StageCheckpointProgresses` + pub fn get_pre_encoded(&self) -> Option<&Vec> { + #[cfg(not(feature = "std"))] + { + None + } + #[cfg(feature = "std")] + ENCODED_STAGE_IDS + .get_or_init(|| { + let mut map = HashMap::with_capacity(Self::ALL.len()); + for stage_id in Self::ALL { + map.insert(stage_id, stage_id.to_string().into_bytes()); + } + map + }) + .get(self) + } } impl core::fmt::Display for StageId { diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs index 5028ffcc88b..160ed34a176 100644 --- a/crates/storage/provider/src/providers/database/provider.rs +++ b/crates/storage/provider/src/providers/database/provider.rs @@ -1642,7 +1642,11 @@ impl BlockBodyIndicesProvider impl StageCheckpointReader for DatabaseProvider { fn get_stage_checkpoint(&self, id: StageId) -> ProviderResult> { - Ok(self.tx.get::(id.to_string())?) + Ok(if let Some(encoded) = id.get_pre_encoded() { + self.tx.get_by_encoded_key::(encoded)? + } else { + self.tx.get::(id.to_string())? + }) } /// Get stage checkpoint progress.