Outline of scalable key value store interface (#13083)

crates/sui-core/src/authority/authority_store.rs

@@ -7,6 +7,7 @@ use std::ops::Not;
 use std::sync::Arc;
 use std::{iter, mem, thread};
 
+use async_trait::async_trait;
 use either::Either;
 use fastcrypto::hash::{HashFunction, MultisetHash, Sha3_256};
 use futures::stream::FuturesUnordered;
@@ -42,6 +43,7 @@ use crate::authority::epoch_start_configuration::{EpochFlag, EpochStartConfigura
 use super::authority_store_tables::LiveObject;
 use super::{authority_store_tables::AuthorityPerpetualTables, *};
 use mysten_common::sync::notify_read::NotifyRead;
+use sui_storage::key_value_store;
 use sui_storage::package_object_cache::PackageObjectCache;
 use sui_types::effects::{TransactionEffects, TransactionEvents};
 use sui_types::gas_coin::TOTAL_SUPPLY_MIST;
@@ -1932,6 +1934,93 @@ impl GetModule for AuthorityStore {
     }
 }
 
+#[async_trait]
+impl key_value_store::TransactionKeyValueStore for AuthorityStore {
+    /// Generic multi_get, allows implementors to get heterogenous values with a single round trip.
+    async fn multi_get(
+        &self,
+        keys: &[key_value_store::Key],
+    ) -> SuiResult<Vec<Option<key_value_store::Value>>> {
+        let mut tx_keys = Vec::new();
+        let mut fx_keys = Vec::new();
+        let mut events_keys = Vec::new();
+
+        for key in keys {
+            match key {
+                key_value_store::Key::Tx(digest) => tx_keys.push(*digest),
+                key_value_store::Key::Fx(digest) => fx_keys.push(*digest),
+                key_value_store::Key::Events(digest) => events_keys.push(*digest),
+            }
+        }
+
+        let tx_results = if !tx_keys.is_empty() {
+            self.multi_get_tx(&tx_keys).await?
+        } else {
+            Vec::new()
+        };
+
+        let fx_results = if !fx_keys.is_empty() {
+            self.multi_get_fx(&fx_keys).await?
+        } else {
+            Vec::new()
+        };
+
+        let events_results = if !events_keys.is_empty() {
+            key_value_store::TransactionKeyValueStore::multi_get_events(self, &events_keys).await?
+        } else {
+            Vec::new()
+        };
+
+        // re-assemble original order
+        let mut tx_iter = tx_results.into_iter();
+        let mut fx_iter = fx_results.into_iter();
+        let mut events_iter = events_results.into_iter();
+
+        let mut results = Vec::new();
+
+        for key in keys {
+            match key {
+                key_value_store::Key::Tx(_) => {
+                    results.push(tx_iter.next().unwrap().map(|tx| tx.into()))
+                }
+                key_value_store::Key::Fx(_) => {
+                    results.push(fx_iter.next().unwrap().map(|fx| fx.into()))
+                }
+                key_value_store::Key::Events(_) => {
+                    results.push(events_iter.next().unwrap().map(|e| e.into()))
+                }
+            }
+        }
+
+        Ok(results)
+    }
+
+    async fn multi_get_tx(
+        &self,
+        keys: &[TransactionDigest],
+    ) -> SuiResult<Vec<Option<Transaction>>> {
+        Ok(self
+            .multi_get_transaction_blocks(keys)?
+            .into_iter()
+            .map(|t| t.map(|t| t.into_inner()))
+            .collect())
+    }
+
+    async fn multi_get_fx(
+        &self,
+        keys: &[TransactionEffectsDigest],
+    ) -> SuiResult<Vec<Option<TransactionEffects>>> {
+        Ok(self.multi_get_effects(keys.iter())?)
+    }
+
+    async fn multi_get_events(
+        &self,
+        keys: &[TransactionEventsDigest],
+    ) -> SuiResult<Vec<Option<TransactionEvents>>> {
+        Ok(self.multi_get_events(keys)?)
+    }
+}
+
 /// A wrapper to make Orphan Rule happy
 pub struct ResolverWrapper<T: ModuleResolver> {
     pub resolver: Arc<T>,

crates/sui-storage/Cargo.toml

@@ -52,3 +52,5 @@ tempfile.workspace = true
 num_cpus.workspace = true
 pretty_assertions.workspace = true
 once_cell.workspace = true
+sui-test-transaction-builder.workspace = true
+sui-types = { workspace = true, features = ["test-utils"] }