hyperlight-guest-tracing 0.12.0

/*
Copyright 2025  The Hyperlight Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
extern crate alloc;

use alloc::sync::Arc;

use spin::Mutex;
use tracing_core::span::{Attributes, Id, Record};
use tracing_core::subscriber::Subscriber;
use tracing_core::{Event, Metadata};

use crate::state::GuestState;

/// The subscriber is used to collect spans and events in the guest.
pub(crate) struct GuestSubscriber {
    /// Internal state that holds the spans and events
    /// Protected by a Mutex for inner mutability
    /// A reference to this state is stored in a static variable
    /// so it can be accessed from the guest tracing API
    state: Arc<Mutex<GuestState>>,
}

impl GuestSubscriber {
    pub(crate) fn new(guest_start_tsc: u64) -> Self {
        Self {
            state: Arc::new(Mutex::new(GuestState::new(guest_start_tsc))),
        }
    }
    pub(crate) fn state(&self) -> &Arc<Mutex<GuestState>> {
        &self.state
    }
}

impl Subscriber for GuestSubscriber {
    fn enabled(&self, _md: &Metadata<'_>) -> bool {
        true
    }

    fn new_span(&self, attrs: &Attributes<'_>) -> Id {
        // We want to protect against re-entrancy issues produced by tracing code that locks
        // the state and then causes an exception that tries to lock the state again.
        //
        // For example:
        // - 1. A span is created, locking the state
        // - 2. An exception occurs while the span is being created (e.g. not enough memory, etc.)
        // - 3. The exception handler uses the tracing API to send the trace data to the host
        // or just create spans/events for logging purposes.
        // - 4. The tracing API tries to lock the state again, causing a deadlock.
        // To avoid this, we use try_lock and if we cannot acquire the lock, we panic to signal
        // the issue.
        let mut state = self
            .state
            .try_lock()
            .expect("guest_tracing: Unable to lock guest tracing state in `new_span`");

        state.new_span(attrs)
    }

    fn record(&self, id: &Id, values: &Record<'_>) {
        // We want to protect against re-entrancy issues produced by tracing code that locks
        // the state and then causes an exception that tries to lock the state again.
        //
        // For example:
        // - 1. A span is created, locking the state
        // - 2. An exception occurs while the span is being created (e.g. not enough memory, etc.)
        // - 3. The exception handler uses the tracing API to send the trace data to the host
        // or just create spans/events for logging purposes.
        // - 4. The tracing API tries to lock the state again, causing a deadlock.
        // To avoid this, we use try_lock and if we cannot acquire the lock, we panic to signal
        // the issue.
        let mut state = self
            .state
            .try_lock()
            .expect("guest_tracing: Unable to lock guest tracing state in `record`");

        state.record(id, values)
    }

    fn event(&self, event: &Event<'_>) {
        // We want to protect against re-entrancy issues produced by tracing code that locks
        // the state and then causes an exception that tries to lock the state again.
        //
        // For example:
        // - 1. A span is created, locking the state
        // - 2. An exception occurs while the span is being created (e.g. not enough memory, etc.)
        // - 3. The exception handler uses the tracing API to send the trace data to the host
        // or just create spans/events for logging purposes.
        // - 4. The tracing API tries to lock the state again, causing a deadlock.
        // To avoid this, we use try_lock and if we cannot acquire the lock, we panic to signal
        // the issue.
        let mut state = self
            .state
            .try_lock()
            .expect("guest_tracing: Unable to lock guest tracing state in `event`");

        state.event(event)
    }

    fn enter(&self, id: &Id) {
        // We want to protect against re-entrancy issues produced by tracing code that locks
        // the state and then causes an exception that tries to lock the state again.
        //
        // For example:
        // - 1. A span is created, locking the state
        // - 2. An exception occurs while the span is being created (e.g. not enough memory, etc.)
        // - 3. The exception handler uses the tracing API to send the trace data to the host
        // or just create spans/events for logging purposes.
        // - 4. The tracing API tries to lock the state again, causing a deadlock.
        // To avoid this, we use try_lock and if we cannot acquire the lock, we panic to signal
        // the issue.
        let mut state = self
            .state
            .try_lock()
            .expect("guest_tracing: Unable to lock guest tracing state in `enter`");

        state.enter(id)
    }

    fn exit(&self, id: &Id) {
        // We want to protect against re-entrancy issues produced by tracing code that locks
        // the state and then causes an exception that tries to lock the state again.
        //
        // For example:
        // - 1. A span is created, locking the state
        // - 2. An exception occurs while the span is being created (e.g. not enough memory, etc.)
        // - 3. The exception handler uses the tracing API to send the trace data to the host
        // or just create spans/events for logging purposes.
        // - 4. The tracing API tries to lock the state again, causing a deadlock.
        // To avoid this, we use try_lock and if we cannot acquire the lock, we panic to signal
        // the issue.
        let mut state = self
            .state
            .try_lock()
            .expect("guest_tracing: Unable to lock guest tracing state in `exit`");

        state.exit(id)
    }

    fn try_close(&self, id: Id) -> bool {
        // We want to protect against re-entrancy issues produced by tracing code that locks
        // the state and then causes an exception that tries to lock the state again.
        //
        // For example:
        // - 1. A span is created, locking the state
        // - 2. An exception occurs while the span is being created (e.g. not enough memory, etc.)
        // - 3. The exception handler uses the tracing API to send the trace data to the host
        // or just create spans/events for logging purposes.
        // - 4. The tracing API tries to lock the state again, causing a deadlock.
        // To avoid this, we use try_lock and if we cannot acquire the lock, we panic to signal
        // the issue.
        let mut state = self
            .state
            .try_lock()
            .expect("guest_tracing: Unable to lock guest tracing state in `try_close`");

        state.try_close(id)
    }

    fn record_follows_from(&self, _span: &Id, _follows: &Id) {
        // no-op: we don't track follows-from relationships
    }
}