Make serializing a FunctionCall return a &[u8] instead of Vec<u8>.

This should save memory allocations, but require that a
FlatBufferBuilder is passed in.

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

Cargo.lock

@@ -16,6 +16,7 @@ anyhow = { version = "1.0.99", default-features = false }
 serde_json = { version = "1.0", default-features = false, features = ["alloc"] }
 hyperlight-common = { workspace = true }
 hyperlight-guest-tracing = { workspace = true, default-features = false }
+flatbuffers = { version= "25.2.10", default-features = false }
 
 [features]
 default = []

src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs

@@ -19,6 +19,7 @@ use alloc::string::ToString;
 use alloc::vec::Vec;
 use core::slice::from_raw_parts;
 
+use flatbuffers::FlatBufferBuilder;
 use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
 use hyperlight_common::flatbuffer_wrappers::function_types::{
     ParameterValue, ReturnType, ReturnValue,
@@ -27,6 +28,7 @@ use hyperlight_common::flatbuffer_wrappers::guest_error::{ErrorCode, GuestError}
 use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData;
 use hyperlight_common::flatbuffer_wrappers::guest_log_level::LogLevel;
 use hyperlight_common::flatbuffer_wrappers::host_function_details::HostFunctionDetails;
+use hyperlight_common::flatbuffer_wrappers::util::estimate_flatbuffer_capacity;
 use hyperlight_common::outb::OutBAction;
 
 use super::handle::GuestHandle;
@@ -92,17 +94,19 @@ impl GuestHandle {
         parameters: Option<Vec<ParameterValue>>,
         return_type: ReturnType,
     ) -> Result<()> {
+        let estimated_capacity =
+            estimate_flatbuffer_capacity(function_name, parameters.as_deref().unwrap_or(&[]));
+
         let host_function_call = FunctionCall::new(
             function_name.to_string(),
             parameters,
             FunctionCallType::Host,
             return_type,
         );
 
-        let host_function_call_buffer: Vec<u8> = host_function_call
-            .try_into()
-            .expect("Unable to serialize host function call");
+        let mut builder = FlatBufferBuilder::with_capacity(estimated_capacity);
 
+        let host_function_call_buffer = host_function_call.encode(&mut builder);
         self.push_shared_output_data(host_function_call_buffer)?;
 
         unsafe {
@@ -155,7 +159,7 @@ impl GuestHandle {
             .try_into()
             .expect("Invalid guest_error_buffer, could not be converted to a Vec<u8>");
 
-        if let Err(e) = self.push_shared_output_data(guest_error_buffer) {
+        if let Err(e) = self.push_shared_output_data(&guest_error_buffer) {
             panic!("Unable to push guest error to shared output data: {:#?}", e);
         }
     }
@@ -184,7 +188,7 @@ impl GuestHandle {
             .try_into()
             .expect("Failed to convert GuestLogData to bytes");
 
-        self.push_shared_output_data(bytes)
+        self.push_shared_output_data(&bytes)
             .expect("Unable to push log data to shared output data");
 
         unsafe {

src/hyperlight_guest/Cargo.toml

@@ -16,7 +16,6 @@ limitations under the License.
 
 use alloc::format;
 use alloc::string::ToString;
-use alloc::vec::Vec;
 use core::any::type_name;
 use core::slice::from_raw_parts_mut;
 
@@ -93,7 +92,7 @@ impl GuestHandle {
 
     /// Pushes the given data onto the shared output data buffer.
     #[hyperlight_guest_tracing::trace_function]
-    pub fn push_shared_output_data(&self, data: Vec<u8>) -> Result<()> {
+    pub fn push_shared_output_data(&self, data: &[u8]) -> Result<()> {
         let peb_ptr = self.peb().unwrap();
         let output_stack_size = unsafe { (*peb_ptr).output_stack.size as usize };
         let output_stack_ptr = unsafe { (*peb_ptr).output_stack.ptr as *mut u8 };
@@ -139,7 +138,7 @@ impl GuestHandle {
 
         // write the actual data
         hyperlight_guest_tracing::trace!("copy data", {
-            odb[stack_ptr_rel as usize..stack_ptr_rel as usize + data.len()].copy_from_slice(&data);
+            odb[stack_ptr_rel as usize..stack_ptr_rel as usize + data.len()].copy_from_slice(data);
         });
 
         // write the offset to the newly written data, to the top of the stack

src/hyperlight_guest/src/guest_handle/host_comm.rs

@@ -98,7 +98,7 @@ fn internal_dispatch_function() -> Result<()> {
         handle.write_error(e.kind, Some(e.message.as_str()));
     })?;
 
-    handle.push_shared_output_data(result_vec)
+    handle.push_shared_output_data(&result_vec)
 }
 
 // This is implemented as a separate function to make sure that epilogue in the internal_dispatch_function is called before the halt()

src/hyperlight_guest/src/guest_handle/io.rs

@@ -15,8 +15,10 @@ limitations under the License.
 */
 
 use criterion::{Criterion, criterion_group, criterion_main};
+use flatbuffers::FlatBufferBuilder;
 use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
 use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType};
+use hyperlight_common::flatbuffer_wrappers::util::estimate_flatbuffer_capacity;
 use hyperlight_host::GuestBinary;
 use hyperlight_host::sandbox::{MultiUseSandbox, SandboxConfiguration, UninitializedSandbox};
 use hyperlight_testing::{c_simple_guest_as_string, simple_guest_as_string};
@@ -157,20 +159,26 @@ fn function_call_serialization_benchmark(c: &mut Criterion) {
         ReturnType::Int,
     );
 
-    let serialized_bytes: Vec<u8> = function_call.clone().try_into().unwrap();
-
     group.bench_function("serialize_function_call", |b| {
-        b.iter_with_setup(
-            || function_call.clone(),
-            |fc| {
-                let _serialized: Vec<u8> = fc.try_into().unwrap();
-            },
-        );
+        b.iter(|| {
+            // We specifically want to include the time to estimate the capacity in this benchmark
+            let estimated_capacity = estimate_flatbuffer_capacity(
+                function_call.function_name.as_str(),
+                function_call.parameters.as_deref().unwrap_or(&[]),
+            );
+            let mut builder = FlatBufferBuilder::with_capacity(estimated_capacity);
+            let serialized: &[u8] = function_call.encode(&mut builder);
+            std::hint::black_box(serialized);
+        });
     });
 
     group.bench_function("deserialize_function_call", |b| {
+        let mut builder = FlatBufferBuilder::new();
+        let bytes = function_call.clone().encode(&mut builder);
+
         b.iter(|| {
-            let _deserialized: FunctionCall = serialized_bytes.as_slice().try_into().unwrap();
+            let deserialized: FunctionCall = bytes.try_into().unwrap();
+            std::hint::black_box(deserialized);
         });
     });
 

src/hyperlight_guest_bin/src/guest_function/call.rs

@@ -23,10 +23,12 @@ use std::path::Path;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
 
+use flatbuffers::FlatBufferBuilder;
 use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
 use hyperlight_common::flatbuffer_wrappers::function_types::{
     ParameterValue, ReturnType, ReturnValue,
 };
+use hyperlight_common::flatbuffer_wrappers::util::estimate_flatbuffer_capacity;
 use tracing::{Span, instrument};
 
 use super::host_funcs::FunctionRegistry;
@@ -44,7 +46,7 @@ use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
 use crate::mem::ptr::RawPtr;
 use crate::mem::shared_mem::HostSharedMemory;
 use crate::metrics::maybe_time_and_emit_guest_call;
-use crate::{HyperlightError, Result, log_then_return};
+use crate::{Result, log_then_return};
 
 /// Global counter for assigning unique IDs to sandboxes
 static SANDBOX_ID_COUNTER: AtomicU64 = AtomicU64::new(0);
@@ -392,20 +394,21 @@ impl MultiUseSandbox {
         args: Vec<ParameterValue>,
     ) -> Result<ReturnValue> {
         let res = (|| {
+            let estimated_capacity = estimate_flatbuffer_capacity(function_name, &args);
+
             let fc = FunctionCall::new(
                 function_name.to_string(),
                 Some(args),
                 FunctionCallType::Guest,
                 return_type,
             );
 
-            let buffer: Vec<u8> = fc.try_into().map_err(|_| {
-                HyperlightError::Error("Failed to serialize FunctionCall".to_string())
-            })?;
+            let mut builder = FlatBufferBuilder::with_capacity(estimated_capacity);
+            let buffer = fc.encode(&mut builder);
 
             self.get_mgr_wrapper_mut()
                 .as_mut()
-                .write_guest_function_call(&buffer)?;
+                .write_guest_function_call(buffer)?;
 
             self.vm.dispatch_call_from_host(
                 self.dispatch_ptr.clone(),