[trace] Improve guest traces to contain top spans

- For each call to `internal_dispatch_function` now we create a span
  manually because instrumenting the function would be redundant because
  of the call to `new_call` that resets tracing state for each new
  function call
- The same logic is used for `generic_init` to enable it for the
  initialise call

Signed-off-by: Doru Blânzeanu <dblnz@pm.me>

Author: dblnz

src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs

@@ -41,7 +41,7 @@ pub enum FunctionCallType {
 }
 
 /// `Functioncall` represents a call to a function in the guest or host.
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct FunctionCall {
     /// The function name
     pub function_name: String,

src/hyperlight_guest_bin/src/guest_function/call.rs

@@ -81,23 +81,25 @@ pub(crate) fn call_guest_function(function_call: FunctionCall) -> Result<Vec<u8>
 pub(crate) fn internal_dispatch_function() {
     // Read the current TSC to report it to the host with the spans/events
     // This helps calculating the timestamps relative to the guest call
-    #[cfg(feature = "trace_guest")]
-    {
+    #[cfg(all(feature = "trace_guest", target_arch = "x86_64"))]
+    let _entered = {
         let guest_start_tsc = hyperlight_guest_tracing::invariant_tsc::read_tsc();
         // Reset the trace state for the new guest function call with the new start TSC
         // This clears any existing spans/events from previous calls ensuring a clean state
         hyperlight_guest_tracing::new_call(guest_start_tsc);
-    }
 
-    let handle = unsafe { GUEST_HANDLE };
+        tracing::span!(tracing::Level::INFO, "internal_dispatch_function").entered()
+    };
 
-    #[cfg(debug_assertions)]
-    log::trace!("internal_dispatch_function");
+    let handle = unsafe { GUEST_HANDLE };
 
     let function_call = handle
         .try_pop_shared_input_data_into::<FunctionCall>()
         .expect("Function call deserialization failed");
 
+    #[cfg(debug_assertions)]
+    tracing::trace!("{:?}", function_call);
+
     let res = call_guest_function(function_call);
 
     match res {
@@ -117,8 +119,20 @@ pub(crate) fn internal_dispatch_function() {
         }
     }
 
-    // Ensure that any tracing output during the call is flushed to
-    // the host, if necessary.
+    // All this tracing logic shall be done right before the call to `hlt` which is done after this
+    // function returns
     #[cfg(all(feature = "trace_guest", target_arch = "x86_64"))]
-    hyperlight_guest_tracing::flush();
+    {
+        // This span captures the internal dispatch function only, without tracing internals.
+        // Close the span before flushing to ensure that the `flush` call is not included in the span
+        // NOTE: This is necessary to avoid closing the span twice. Flush closes all the open
+        // spans, when preparing to close a guest function call context.
+        // It is not mandatory, though, but avoids a warning on the host that alerts a spans
+        // that has not been opened but is being closed.
+        _entered.exit();
+
+        // Ensure that any tracing output during the call is flushed to
+        // the host, if necessary.
+        hyperlight_guest_tracing::flush();
+    }
 }

src/hyperlight_guest_bin/src/lib.rs

@@ -226,6 +226,12 @@ pub(crate) extern "C" fn generic_init(
         hyperlight_guest_tracing::init_guest_tracing(guest_start_tsc);
     }
 
+    // Open a span to partly capture the initialization of the guest.
+    // This is done here because the tracing subscriber is initialized and the guest is in a
+    // well-known state
+    #[cfg(all(feature = "trace_guest", target_arch = "x86_64"))]
+    let _entered = tracing::span!(tracing::Level::INFO, "generic_init").entered();
+
     #[cfg(feature = "macros")]
     for registration in __private::GUEST_FUNCTION_INIT {
         registration();
@@ -235,10 +241,22 @@ pub(crate) extern "C" fn generic_init(
         hyperlight_main();
     }
 
-    // Ensure that any tracing output from the initialisation phase is
-    // flushed to the host, if necessary.
+    // All this tracing logic shall be done right before the call to `hlt` which is done after this
+    // function returns
     #[cfg(all(feature = "trace_guest", target_arch = "x86_64"))]
-    hyperlight_guest_tracing::flush();
+    {
+        // This span captures the internal dispatch function only, without tracing internals.
+        // Close the span before flushing to ensure that the `flush` call is not included in the span
+        // NOTE: This is necessary to avoid closing the span twice. Flush closes all the open
+        // spans, when preparing to close a guest function call context.
+        // It is not mandatory, though, but avoids a warning on the host that alerts a spans
+        // that has not been opened but is being closed.
+        _entered.exit();
+
+        // Ensure that any tracing output from the initialisation phase is
+        // flushed to the host, if necessary.
+        hyperlight_guest_tracing::flush();
+    }
 
     dispatch_function as usize as u64
 }