[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

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.2.10", 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::{Span, instrument};
 
 /// Halt the execution of the guest and returns control to the host.
 #[inline(never)]
+#[instrument(skip_all, parent = Span::current(), level= "Trace")]
 pub fn halt() {
     #[cfg(feature = "trace_guest")]
     {
@@ -93,6 +95,7 @@ pub unsafe fn abort_with_code_and_message(code: &[u8], message_ptr: *const c_cha
 }
 
 /// OUT bytes to the host through multiple exits.
+#[instrument(skip_all, parent = Span::current(), level= "Trace")]
 pub(crate) fn outb(port: u16, data: &[u8]) {
     // Ensure all tracing data is flushed before sending OUT bytes
     unsafe {
@@ -111,6 +114,7 @@ pub(crate) fn outb(port: u16, data: &[u8]) {
 }
 
 /// OUT function for sending a 32-bit value to the host.
+#[instrument(skip_all, parent = Span::current(), 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::{Span, 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, parent = Span::current(), 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 };
@@ -65,6 +67,7 @@ impl GuestHandle {
     ///
     /// When calling `call_host_function<T>`, this function is called
     /// internally to get the return value.
+    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     pub fn get_host_return_value<T: TryFrom<ReturnValue>>(&self) -> Result<T> {
         let return_value = self
             .try_pop_shared_input_data_into::<ReturnValue>()
@@ -85,6 +88,7 @@ impl GuestHandle {
     ///
     /// Note: The function return value must be obtained by calling
     /// `get_host_return_value`.
+    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     pub fn call_host_function_without_returning_result(
         &self,
         function_name: &str,
@@ -118,6 +122,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, parent = Span::current(), level= "Trace")]
     pub fn call_host_function<T: TryFrom<ReturnValue>>(
         &self,
         function_name: &str,
@@ -128,6 +133,7 @@ impl GuestHandle {
         self.get_host_return_value::<T>()
     }
 
+    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     pub fn get_host_function_details(&self) -> HostFunctionDetails {
         let peb_ptr = self.peb().unwrap();
         let host_function_details_buffer =
@@ -144,6 +150,7 @@ impl GuestHandle {
     }
 
     /// Write an error to the shared output data buffer.
+    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     pub fn write_error(&self, error_code: ErrorCode, message: Option<&str>) {
         let guest_error: GuestError = GuestError::new(
             error_code,
@@ -159,6 +166,7 @@ impl GuestHandle {
     }
 
     /// Log a message with the specified log level, source, caller, source file, and line number.
+    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     pub fn log_message(
         &self,
         log_level: LogLevel,
@@ -168,24 +176,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

@@ -27,6 +27,7 @@ hyperlight-guest-tracing = { workspace = true, default-features = false }
 buddy_system_allocator = "0.11.0"
 log = { version = "0.4", default-features = false }
 spin = "0.10.0"
+tracing = { version = "0.1.41", default-features = false, features = ["attributes"] }
 
 [lints]
 workspace = true

src/hyperlight_guest_bin/Cargo.toml

@@ -22,11 +22,13 @@ use hyperlight_common::flatbuffer_wrappers::function_types::ParameterType;
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 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 {
@@ -80,6 +82,7 @@ pub(crate) fn call_guest_function(function_call: FunctionCall) -> Result<Vec<u8>
 // and we will leak memory as the epilogue will not be called as halt() is not going to return.
 #[unsafe(no_mangle)]
 #[inline(never)]
+#[instrument(skip_all, parent = Span::current(), level= "Trace")]
 fn internal_dispatch_function() -> Result<()> {
     let handle = unsafe { GUEST_HANDLE };
 
@@ -100,6 +103,7 @@ fn internal_dispatch_function() -> Result<()> {
 // 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

@@ -16,7 +16,9 @@ limitations under the License.
 #![no_std]
 use heapless as hl;
 #[cfg(feature = "trace")]
-pub use trace::{TraceBatchInfo, end_trace, guest_trace_info, init_guest_tracing};
+pub use trace::{
+    TraceBatchInfo, end_trace, guest_trace_info, init_guest_tracing, is_trace_enabled,
+};
 
 /// Module for checking invariant TSC support and reading the timestamp counter
 pub mod invariant_tsc {
@@ -546,4 +548,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()
+            .and_then(|w| Some(w.upgrade().is_some()))
+            .unwrap_or(false)
+    }
 }