1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
//! The JIT (Just-In-Time) Profiling API provides functionality to report information about
//! just-in-time generated code that can be used by performance tools. The [Jit] Rust structure is a
//! high-level view of a subset of the full functionality available. See the [JIT Profiling API] for
//! more information.
//!
//! [JIT Profiling API]:
///     https://www.intel.com/content/www/us/en/develop/documentation/vtune-help/top/api-support/jit-profiling-api.html
use anyhow::Context;
use std::{ffi::CString, os, ptr};

/// Register JIT-compiled methods with a performance tool (e.g., VTune). This structure assumes
/// single-threaded access; if your program may be multi-threaded, make sure to guard multi-threaded
/// access with a mutex.
#[derive(Default)]
pub struct Jit {
    shutdown_complete: bool,
}

impl Jit {
    /// Returns a new `MethodId` for use in `MethodLoad` events.
    #[allow(clippy::unused_self)]
    pub fn get_method_id(&self) -> MethodId {
        MethodId(unsafe { ittapi_sys::iJIT_GetNewMethodID() })
    }

    /// Notifies any `EventType` to VTune.
    ///
    /// # Errors
    ///
    /// May fail if the underlying call to the ITT library fails.
    #[allow(clippy::unused_self)]
    pub fn notify_event(&self, mut event: EventType) -> anyhow::Result<()> {
        let tag = event.tag();
        let data = event.data();
        log::trace!("notify_event: tag={:?}", tag);
        let res = unsafe { ittapi_sys::iJIT_NotifyEvent(tag, data) };
        if res == 1 {
            Ok(())
        } else {
            anyhow::bail!("error when notifying event")
        }
    }

    // High-level helpers.

    /// Notifies VTune that a new function described by the `MethodLoadBuilder` has been jitted.
    ///
    /// # Errors
    ///
    /// May fail if the builder has invalid data or if the ITT library fails to notify the method
    /// load event.
    pub fn load_method(&self, builder: MethodLoadBuilder) -> anyhow::Result<()> {
        let method_id = self.get_method_id();
        let method_load = builder.build(method_id)?;
        self.notify_event(method_load)
    }

    /// Notifies VTune that profiling is being shut down.
    ///
    /// # Errors
    ///
    /// May fail if the ITT library fails to notify the shutdown event.
    pub fn shutdown(&mut self) -> anyhow::Result<()> {
        let res = self.notify_event(EventType::Shutdown);
        if res.is_ok() {
            self.shutdown_complete = true;
        }
        res
    }
}

impl Drop for Jit {
    fn drop(&mut self) {
        if !self.shutdown_complete {
            // There's not much we can do when an error happens here.
            if let Err(err) = self.shutdown() {
                log::error!("Error when shutting down VTune: {}", err);
            }
        }
    }
}

/// Type of event to be dispatched through ittapi's JIT event API.
pub enum EventType {
    /// Send this event after a JITted method has been loaded into memory, and possibly JIT
    /// compiled, but before the code is executed.
    MethodLoadFinished(MethodLoad),

    /// Send this notification to terminate profiling.
    Shutdown,
}

impl EventType {
    /// Returns the C event type to be used when notifying this event to VTune.
    fn tag(&self) -> ittapi_sys::iJIT_jvm_event {
        match self {
            EventType::MethodLoadFinished(_) => {
                ittapi_sys::iJIT_jvm_event_iJVM_EVENT_TYPE_METHOD_INLINE_LOAD_FINISHED
            }
            EventType::Shutdown => ittapi_sys::iJIT_jvm_event_iJVM_EVENT_TYPE_SHUTDOWN,
        }
    }

    /// Returns a raw C pointer to the event-specific data that must be used when notifying this
    /// event to VTune.
    fn data(&mut self) -> *mut os::raw::c_void {
        match self {
            EventType::MethodLoadFinished(method_load) => ptr::addr_of_mut!(method_load.0).cast(),
            EventType::Shutdown => ptr::null_mut(),
        }
    }
}

/// Newtype wrapper for a method id returned by ittapi's `iJIT_GetNewMethodID`, as returned by
/// `VtuneState::get_method_id` in the high-level API.
#[derive(Clone, Copy)]
pub struct MethodId(u32);

/// Newtype wrapper for a JIT method load.
pub struct MethodLoad(ittapi_sys::_iJIT_Method_Load);

/// Multi-step constructor using the builder pattern for a `MethodLoad` event.
pub struct MethodLoadBuilder {
    method_name: String,
    addr: *const u8,
    len: usize,
    class_file_name: Option<String>,
    source_file_name: Option<String>,
}

impl MethodLoadBuilder {
    /// Creates a new `MethodLoadBuilder` from scratch.
    ///
    /// `addr` is the pointer to the start of the code region, `len` is the size of this code
    /// region in bytes.
    pub fn new(method_name: String, addr: *const u8, len: usize) -> Self {
        Self {
            method_name,
            addr,
            len,
            class_file_name: None,
            source_file_name: None,
        }
    }

    /// Attache a class file.
    pub fn class_file_name(mut self, class_file_name: String) -> Self {
        self.class_file_name = Some(class_file_name);
        self
    }

    /// Attach a source file.
    pub fn source_file_name(mut self, source_file_name: String) -> Self {
        self.source_file_name = Some(source_file_name);
        self
    }

    /// Build a "method load" event for the given `method_id`.
    ///
    /// # Errors
    ///
    /// May fail if the various names passed to this builder are not valid C strings.
    pub fn build(self, method_id: MethodId) -> anyhow::Result<EventType> {
        Ok(EventType::MethodLoadFinished(MethodLoad(
            ittapi_sys::_iJIT_Method_Load {
                method_id: method_id.0,
                method_name: CString::new(self.method_name)
                    .context("CString::new failed")?
                    .into_raw(),
                method_load_address: self.addr as *mut os::raw::c_void,
                method_size: self.len.try_into().expect("cannot fit length into 32 bits"),
                line_number_size: 0,
                line_number_table: ptr::null_mut(),
                class_id: 0, // Field officially obsolete in Intel's doc.
                class_file_name: CString::new(
                    self.class_file_name
                        .as_deref()
                        .unwrap_or("<unknown class file name>"),
                )
                .context("CString::new failed")?
                .into_raw(),
                source_file_name: CString::new(
                    self.source_file_name
                        .as_deref()
                        .unwrap_or("<unknown source file name>"),
                )
                .context("CString::new failed")?
                .into_raw(),
            },
        )))
    }
}