diff --git a/storage-bigtable/src/bigtable.rs b/storage-bigtable/src/bigtable.rs index 7cf864e3f1..6e6190da98 100644 --- a/storage-bigtable/src/bigtable.rs +++ b/storage-bigtable/src/bigtable.rs @@ -680,6 +680,32 @@ impl) -> InterceptedRequestResult> BigTable { deserialize_protobuf_or_bincode_cell_data(&row_data, table, key) } + #[allow(clippy::needless_lifetimes)] + pub async fn get_protobuf_or_bincode_cells<'a, B, P>( + &mut self, + table: &'a str, + row_keys: impl IntoIterator, + ) -> Result)> + 'a> + where + B: serde::de::DeserializeOwned, + P: prost::Message + Default, + { + Ok(self + .get_multi_row_data( + table, + row_keys.into_iter().collect::>().as_slice(), + ) + .await? + .into_iter() + .map(|(key, row_data)| { + let key_str = key.to_string(); + ( + key, + deserialize_protobuf_or_bincode_cell_data(&row_data, table, key_str).unwrap(), + ) + })) + } + pub async fn put_bincode_cells( &mut self, table: &str, diff --git a/storage-bigtable/src/lib.rs b/storage-bigtable/src/lib.rs index 02c27c4167..5f09e7abaa 100644 --- a/storage-bigtable/src/lib.rs +++ b/storage-bigtable/src/lib.rs @@ -1,5 +1,6 @@ #![allow(clippy::integer_arithmetic)] use { + crate::bigtable::RowKey, log::*, serde::{Deserialize, Serialize}, solana_metrics::inc_new_counter_debug, @@ -393,6 +394,37 @@ impl LedgerStorage { Ok(blocks.into_iter().filter_map(|s| key_to_slot(&s)).collect()) } + // Fetches and gets a vector of confirmed blocks via a multirow fetch + #[allow(clippy::needless_lifetimes)] + pub async fn get_confirmed_blocks_with_data<'a>( + &self, + slots: &'a [Slot], + ) -> Result + 'a> { + debug!( + "LedgerStorage::get_confirmed_blocks_with_data request received: {:?}", + slots + ); + inc_new_counter_debug!("storage-bigtable-query", 1); + let mut bigtable = self.connection.client(); + let row_keys = slots.iter().copied().map(slot_to_blocks_key); + let data = bigtable + .get_protobuf_or_bincode_cells("blocks", row_keys) + .await? + .filter_map( + |(row_key, block_cell_data): ( + RowKey, + bigtable::CellData, + )| { + let block = match block_cell_data { + bigtable::CellData::Bincode(block) => block.into(), + bigtable::CellData::Protobuf(block) => block.try_into().ok()?, + }; + Some((key_to_slot(&row_key).unwrap(), block)) + }, + ); + Ok(data) + } + /// Fetch the confirmed block from the desired slot pub async fn get_confirmed_block( &self,