[trace-guest] update guest crates to add trace instrumentation

- conditionally handle logs either through tracing or the dedicated VM exit
  based on whether tracing is initialized on the guest
- modify `test_with_small_stack_and_heap` to 18kB because the `#[intrument]`
  attributes use more stack.

Signed-off-by: Doru Blânzeanu <[email protected]>

Author: dblnz

Cargo.lock

@@ -17,6 +17,7 @@ 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.9.23", default-features = false }
+tracing = { version = "0.1.41", default-features = false, features = ["attributes"] }
 
 [features]
 default = []

src/hyperlight_guest/Cargo.toml

@@ -18,9 +18,11 @@ use core::arch::asm;
 use core::ffi::{CStr, c_char};
 
 use hyperlight_common::outb::OutBAction;
+use tracing::instrument;
 
 /// Halt the execution of the guest and returns control to the host.
 #[inline(never)]
+#[instrument(skip_all, level = "Trace")]
 pub fn halt() {
     #[cfg(feature = "trace_guest")]
     {
@@ -121,6 +123,7 @@ pub(crate) fn outb(port: u16, data: &[u8]) {
 }
 
 /// OUT function for sending a 32-bit value to the host.
+#[instrument(skip_all, level = "Trace")]
 pub(crate) unsafe fn out32(port: u16, val: u32) {
     #[cfg(feature = "trace_guest")]
     {

src/hyperlight_guest/src/exit.rs

@@ -30,13 +30,15 @@ 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 tracing::instrument;
 
 use super::handle::GuestHandle;
 use crate::error::{HyperlightGuestError, Result};
 use crate::exit::out32;
 
 impl GuestHandle {
     /// Get user memory region as bytes.
+    #[instrument(skip_all, level = "Trace")]
     pub fn read_n_bytes_from_user_memory(&self, num: u64) -> Result<Vec<u8>> {
         let peb_ptr = self.peb().unwrap();
         let user_memory_region_ptr = unsafe { (*peb_ptr).init_data.ptr as *mut u8 };
@@ -124,6 +126,7 @@ impl GuestHandle {
     /// sends it to the host, and then retrieves the return value.
     ///
     /// The return value is deserialized into the specified type `T`.
+    #[instrument(skip_all, level = "Trace")]
     pub fn call_host_function<T: TryFrom<ReturnValue>>(
         &self,
         function_name: &str,
@@ -134,6 +137,7 @@ impl GuestHandle {
         self.get_host_return_value::<T>()
     }
 
+    #[instrument(skip_all, level = "Trace")]
     pub fn get_host_function_details(&self) -> HostFunctionDetails {
         let peb_ptr = self.peb().unwrap();
         let host_function_details_buffer =
@@ -159,24 +163,46 @@ impl GuestHandle {
         source_file: &str,
         line: u32,
     ) {
-        let guest_log_data = GuestLogData::new(
-            message.to_string(),
-            source.to_string(),
-            log_level,
-            caller.to_string(),
-            source_file.to_string(),
-            line,
-        );
-
-        let bytes: Vec<u8> = guest_log_data
-            .try_into()
-            .expect("Failed to convert GuestLogData to bytes");
-
-        self.push_shared_output_data(&bytes)
-            .expect("Unable to push log data to shared output data");
-
-        unsafe {
-            out32(OutBAction::Log as u16, 0);
+        // Closure to send log message to host
+        let send_to_host = || {
+            let guest_log_data = GuestLogData::new(
+                message.to_string(),
+                source.to_string(),
+                log_level,
+                caller.to_string(),
+                source_file.to_string(),
+                line,
+            );
+
+            let bytes: Vec<u8> = guest_log_data
+                .try_into()
+                .expect("Failed to convert GuestLogData to bytes");
+
+            self.push_shared_output_data(&bytes)
+                .expect("Unable to push log data to shared output data");
+
+            unsafe {
+                out32(OutBAction::Log as u16, 0);
+            }
+        };
+
+        #[cfg(feature = "trace_guest")]
+        if hyperlight_guest_tracing::is_trace_enabled() {
+            // If the "trace_guest" feature is enabled and tracing is initialized, log using tracing
+            tracing::trace!(
+                event = message,
+                level = ?log_level,
+                code.filepath = source,
+                caller = caller,
+                source_file = source_file,
+                code.lineno = line,
+            );
+        } else {
+            send_to_host();
+        }
+        #[cfg(not(feature = "trace_guest"))]
+        {
+            send_to_host();
         }
     }
 }

src/hyperlight_guest/src/guest_handle/host_comm.rs

@@ -20,12 +20,14 @@ use core::any::type_name;
 use core::slice::from_raw_parts_mut;
 
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
+use tracing::{Span, instrument};
 
 use super::handle::GuestHandle;
 use crate::error::{HyperlightGuestError, Result};
 
 impl GuestHandle {
     /// Pops the top element from the shared input data buffer and returns it as a T
+    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     pub fn try_pop_shared_input_data_into<T>(&self) -> Result<T>
     where
         T: for<'a> TryFrom<&'a [u8]>,
@@ -87,6 +89,7 @@ impl GuestHandle {
     }
 
     /// Pushes the given data onto the shared output data buffer.
+    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     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 };

src/hyperlight_guest/src/guest_handle/io.rs

@@ -28,7 +28,7 @@ buddy_system_allocator = "0.11.0"
 log = { version = "0.4", default-features = false }
 spin = "0.10.0"
 flatbuffers = { version = "25.2.10", default-features = false }
-
+tracing = { version = "0.1.41", default-features = false, features = ["attributes"] }
 
 [lints]
 workspace = true

src/hyperlight_guest_bin/Cargo.toml

@@ -23,11 +23,13 @@ use hyperlight_common::flatbuffer_wrappers::function_types::{FunctionCallResult,
 use hyperlight_common::flatbuffer_wrappers::guest_error::{ErrorCode, GuestError};
 use hyperlight_guest::error::{HyperlightGuestError, Result};
 use hyperlight_guest::exit::halt;
+use tracing::{Span, instrument};
 
 use crate::{GUEST_HANDLE, REGISTERED_GUEST_FUNCTIONS};
 
 type GuestFunc = fn(&FunctionCall) -> Result<Vec<u8>>;
 
+#[instrument(skip_all, parent = Span::current(), level= "Trace")]
 pub(crate) fn call_guest_function(function_call: FunctionCall) -> Result<Vec<u8>> {
     // Validate this is a Guest Function Call
     if function_call.function_call_type() != FunctionCallType::Guest {
@@ -83,6 +85,7 @@ pub(crate) fn call_guest_function(function_call: FunctionCall) -> Result<Vec<u8>
 // This function may panic, as we have no other ways of dealing with errors at this level
 #[unsafe(no_mangle)]
 #[inline(never)]
+#[instrument(skip_all, parent = Span::current(), level= "Trace")]
 fn internal_dispatch_function() {
     let handle = unsafe { GUEST_HANDLE };
 
@@ -116,6 +119,7 @@ fn internal_dispatch_function() {
 // This is implemented as a separate function to make sure that epilogue in the internal_dispatch_function is called before the halt()
 // which if it were included in the internal_dispatch_function cause the epilogue to not be called because the halt() would not return
 // when running in the hypervisor.
+#[instrument(skip_all, parent = Span::current(), level= "Trace")]
 pub(crate) extern "C" fn dispatch_function() {
     // The hyperlight host likes to use one partition and reset it in
     // various ways; if that has happened, there might stale TLB

src/hyperlight_guest_bin/src/guest_function/call.rs

@@ -17,6 +17,8 @@ limitations under the License.
 use alloc::alloc::Layout;
 use core::arch::asm;
 
+use tracing::{Span, instrument};
+
 use crate::OS_PAGE_SIZE;
 
 /// Convert a physical address in main memory to a virtual address
@@ -61,6 +63,7 @@ struct MapResponse {
 ///   as such do not use concurrently with any other page table operations
 /// - TLB invalidation is not performed,
 ///   if previously-unmapped ranges are not being mapped, TLB invalidation may need to be performed afterwards.
+#[instrument(skip_all, parent = Span::current(), level= "Trace")]
 pub unsafe fn map_region(phys_base: u64, virt_base: *mut u8, len: u64) {
     let mut pml4_base: u64;
     unsafe {

src/hyperlight_guest_bin/src/paging.rs

@@ -38,7 +38,8 @@ mod visitor;
 pub use state::TraceBatchInfo;
 #[cfg(feature = "trace")]
 pub use trace::{
-    clean_trace_state, end_trace, guest_trace_info, init_guest_tracing, set_start_tsc,
+    clean_trace_state, end_trace, guest_trace_info, init_guest_tracing, is_trace_enabled,
+    set_start_tsc,
 };
 
 /// Maximum number of spans that the guest can store
@@ -212,4 +213,12 @@ mod trace {
         }
         res
     }
+
+    /// Returns true if tracing is enabled (the guest tracing state is initialized).
+    pub fn is_trace_enabled() -> bool {
+        GUEST_STATE
+            .get()
+            .map(|w| w.upgrade().is_some())
+            .unwrap_or(false)
+    }
 }

src/hyperlight_guest_tracing/src/lib.rs

@@ -602,7 +602,7 @@ mod tests {
     fn test_with_small_stack_and_heap() {
         let mut cfg = SandboxConfiguration::default();
         cfg.set_heap_size(20 * 1024);
-        cfg.set_stack_size(16 * 1024);
+        cfg.set_stack_size(18 * 1024);
 
         let mut sbox1: MultiUseSandbox = {
             let path = simple_guest_as_string().unwrap();