Reduced files changed, fixed server test.

There is a public datatpe from lookup.rs called Data, which is used externally only to pass data to the lookup module.  This type is used as a Vec<(Bytes, Bytes)> everywhere.  This CL changes Data from HashMap to Vec<(Bytes, Bytes)> for use in external tests.

oak_functions_sdk/tests/integration_test.rs

@@ -14,20 +14,17 @@
 // limitations under the License.
 //
 
-use alloc::{vec, vec::Vec};
 use std::{path::PathBuf, sync::Arc};
 
 use lazy_static::lazy_static;
 use oak_functions_abi::{Request, Response};
 use oak_functions_service::{
     logger::StandaloneLogger,
-    lookup::LookupDataManager,
+    lookup::{Data, LookupDataManager},
     wasm::{api::StdWasmApiFactory, WasmHandler},
     Handler,
 };
 
-extern crate alloc;
-
 lazy_static! {
     static ref PATH_TO_MODULES: PathBuf = {
         // WORKSPACE_ROOT is set in .cargo/config.toml.
@@ -54,7 +51,8 @@ lazy_static! {
 #[tokio::test]
 async fn test_read_write() {
     let logger = Arc::new(StandaloneLogger);
-    let lookup_data_manager = Arc::new(LookupDataManager::for_test(Vec::default(), logger.clone()));
+    let lookup_data_manager =
+        Arc::new(LookupDataManager::for_test(Data::default(), logger.clone()));
     let api_factory = StdWasmApiFactory {
         lookup_data_manager,
     };
@@ -77,7 +75,8 @@ async fn test_read_write() {
 #[tokio::test]
 async fn test_double_read() {
     let logger = Arc::new(StandaloneLogger);
-    let lookup_data_manager = Arc::new(LookupDataManager::for_test(Vec::default(), logger.clone()));
+    let lookup_data_manager =
+        Arc::new(LookupDataManager::for_test(Data::default(), logger.clone()));
     let api_factory = StdWasmApiFactory {
         lookup_data_manager,
     };
@@ -100,7 +99,8 @@ async fn test_double_read() {
 #[tokio::test]
 async fn test_double_write() {
     let logger = Arc::new(StandaloneLogger);
-    let lookup_data_manager = Arc::new(LookupDataManager::for_test(Vec::default(), logger.clone()));
+    let lookup_data_manager =
+        Arc::new(LookupDataManager::for_test(Data::default(), logger.clone()));
     let api_factory = StdWasmApiFactory {
         lookup_data_manager,
     };
@@ -123,7 +123,8 @@ async fn test_double_write() {
 #[tokio::test]
 async fn test_write_log() {
     let logger = Arc::new(StandaloneLogger);
-    let lookup_data_manager = Arc::new(LookupDataManager::for_test(Vec::default(), logger.clone()));
+    let lookup_data_manager =
+        Arc::new(LookupDataManager::for_test(Data::default(), logger.clone()));
     let api_factory = StdWasmApiFactory {
         lookup_data_manager,
     };
@@ -145,10 +146,10 @@ async fn test_write_log() {
 
 #[tokio::test]
 async fn test_storage_get_item() {
-    let entries = vec![(
+    let entries = Data::from_iter([(
         b"StorageGet".to_vec().into(),
         b"StorageGetResponse".to_vec().into(),
-    )];
+    )]);
 
     let logger = Arc::new(StandaloneLogger);
     let lookup_data_manager = Arc::new(LookupDataManager::for_test(entries, logger.clone()));
@@ -174,7 +175,7 @@ async fn test_storage_get_item() {
 #[tokio::test]
 async fn test_storage_get_item_not_found() {
     // empty lookup data, no key will be found
-    let entries = Vec::new();
+    let entries = Data::new();
 
     let logger = Arc::new(StandaloneLogger);
     let lookup_data_manager = Arc::new(LookupDataManager::for_test(entries, logger.clone()));
@@ -201,7 +202,7 @@ async fn test_storage_get_item_not_found() {
 #[ignore]
 async fn test_storage_get_item_huge_key() {
     let bytes: Vec<u8> = vec![42u8; 1 << 20];
-    let entries = Vec::from_iter([(bytes.clone().into(), bytes.clone().into())]);
+    let entries = Data::from_iter([(bytes.clone().into(), bytes.clone().into())]);
 
     let logger = Arc::new(StandaloneLogger);
     let lookup_data_manager = Arc::new(LookupDataManager::for_test(entries, logger.clone()));
@@ -230,7 +231,8 @@ async fn test_echo() {
     let logger = Arc::new(StandaloneLogger);
     let message_to_echo = "ECHO";
 
-    let lookup_data_manager = Arc::new(LookupDataManager::for_test(Vec::default(), logger.clone()));
+    let lookup_data_manager =
+        Arc::new(LookupDataManager::for_test(Data::default(), logger.clone()));
     let api_factory = StdWasmApiFactory {
         lookup_data_manager,
     };
@@ -259,7 +261,8 @@ async fn test_blackhole() {
     let logger = Arc::new(StandaloneLogger);
     let message_to_blackhole = "BLACKHOLE";
 
-    let lookup_data_manager = Arc::new(LookupDataManager::for_test(Vec::default(), logger.clone()));
+    let lookup_data_manager =
+        Arc::new(LookupDataManager::for_test(Data::default(), logger.clone()));
     let api_factory = StdWasmApiFactory {
         lookup_data_manager,
     };
@@ -288,7 +291,8 @@ async fn test_huge_response() {
 
     let logger = Arc::new(StandaloneLogger);
 
-    let lookup_data_manager = Arc::new(LookupDataManager::for_test(Vec::default(), logger.clone()));
+    let lookup_data_manager =
+        Arc::new(LookupDataManager::for_test(Data::default(), logger.clone()));
     let api_factory = StdWasmApiFactory {
         lookup_data_manager,
     };

oak_functions_service/src/lookup.rs

@@ -27,9 +27,12 @@ use spinning_top::{RwSpinlock, Spinlock};
 
 use crate::{logger::OakLogger, lookup_htbl::LookupHtbl};
 
+// This is used to pass data to the lookup data module for tests.
+pub type Data = Vec<(Bytes, Bytes)>;
+
 // Data maintains the invariant on lookup data to have [at most one
 // value](https://github.com/project-oak/oak/tree/main/oak/oak_functions_service/README.md#invariant-at-most-one-value)
-pub type Data = LookupHtbl;
+type DataInternal = LookupHtbl;
 
 #[derive(Default)]
 enum BuilderState {
@@ -41,13 +44,13 @@ enum BuilderState {
 // Incrementally build next lookup data keeping track of the state.
 #[derive(Default)]
 struct DataBuilder {
-    data: Data,
+    data: DataInternal,
     state: BuilderState,
 }
 
 impl DataBuilder {
     /// Build data from the builder and set the builder back to the initial state.
-    fn build(&mut self) -> Data {
+    fn build(&mut self) -> DataInternal {
         self.state = BuilderState::Empty;
         core::mem::take(&mut self.data)
     }
@@ -84,7 +87,7 @@ impl DataBuilder {
 ///
 /// In the future we may replace both the mutex and the hash map with something like RCU.
 pub struct LookupDataManager {
-    data: RwSpinlock<Arc<Data>>,
+    data: RwSpinlock<Arc<DataInternal>>,
     // Behind a lock, because we have multiple references to LookupDataManager and need to mutate
     // data builder.
     data_builder: Spinlock<DataBuilder>,
@@ -95,7 +98,7 @@ impl LookupDataManager {
     /// Creates a new instance with empty backing data.
     pub fn new_empty(logger: Arc<dyn OakLogger>) -> Self {
         Self {
-            data: RwSpinlock::new(Arc::new(Data::default())),
+            data: RwSpinlock::new(Arc::new(DataInternal::default())),
             // Incrementally builds the backing data that will be used by new `LookupData`
             // instances when finished.
             data_builder: Spinlock::new(DataBuilder::default()),
@@ -176,12 +179,12 @@ impl LookupDataManager {
 /// Provides access to shared lookup data.
 #[derive(Clone)]
 pub struct LookupData {
-    data: Arc<Data>,
+    data: Arc<DataInternal>,
     logger: Arc<dyn OakLogger>,
 }
 
 impl LookupData {
-    fn new(data: Arc<Data>, logger: Arc<dyn OakLogger>) -> Self {
+    fn new(data: Arc<DataInternal>, logger: Arc<dyn OakLogger>) -> Self {
         Self { data, logger }
     }
 

oak_functions_service/src/wasm/tests.rs

@@ -28,7 +28,7 @@ use super::{
 };
 use crate::{
     logger::StandaloneLogger,
-    lookup::LookupDataManager,
+    lookup::{Data, LookupDataManager},
     wasm::{AbiPointer, AbiPointerOffset},
     Handler,
 };
@@ -150,7 +150,8 @@ struct TestState {
 
 fn create_test_state() -> TestState {
     let logger = Arc::new(StandaloneLogger);
-    let lookup_data_manager = Arc::new(LookupDataManager::for_test(Vec::new(), logger.clone()));
+    let lookup_data_manager =
+        Arc::new(LookupDataManager::for_test(Data::default(), logger.clone()));
     let api_factory = Arc::new(StdWasmApiFactory {
         lookup_data_manager: lookup_data_manager.clone(),
     });

proto/oak_functions/service/oak_functions.proto

@@ -116,8 +116,6 @@ message AbortNextLookupDataResponse {}
 
 message Empty {}
 
-// The LookupHtbl implementation assumes this is called up-front, and that no
-// duplicate keys exist.
 message ReserveRequest {
   uint64 additional_entries = 1;
 }