perf(db): do not heap-allocate the stage key per query (paradigmxyz#18284)

Author: hai-rise

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<HashMap<StageId, Vec<u8>>> = 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<u8>> {
+        #[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 {

crates/storage/provider/src/providers/database/provider.rs

@@ -1642,7 +1642,11 @@ impl<TX: DbTx + 'static, N: NodeTypesForProvider> BlockBodyIndicesProvider
 
 impl<TX: DbTx, N: NodeTypes> StageCheckpointReader for DatabaseProvider<TX, N> {
     fn get_stage_checkpoint(&self, id: StageId) -> ProviderResult<Option<StageCheckpoint>> {
-        Ok(self.tx.get::<tables::StageCheckpoints>(id.to_string())?)
+        Ok(if let Some(encoded) = id.get_pre_encoded() {
+            self.tx.get_by_encoded_key::<tables::StageCheckpoints>(encoded)?
+        } else {
+            self.tx.get::<tables::StageCheckpoints>(id.to_string())?
+        })
     }
 
     /// Get stage checkpoint progress.