use crate::{JitError, JitResult};
use indexmap::IndexMap;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
#[derive(Debug)]
pub struct ProfilerManager {
sessions: Arc<Mutex<IndexMap<String, ProfilingSession>>>,
counters: Arc<Mutex<PerformanceCounters>>,
config: ProfilerConfig,
external_profilers: Vec<Box<dyn ExternalProfiler>>,
sampling_profiler: Option<SamplingProfiler>,
stats: Arc<Mutex<ProfilerStats>>,
session_counter: Arc<Mutex<u64>>,
}
#[derive(Debug, Clone)]
pub struct ProfilingSession {
pub id: String,
pub name: String,
pub start_time: Instant,
pub duration: Option<Duration>,
pub events: Vec<PerformanceEvent>,
pub call_stacks: Vec<CallStack>,
pub memory_events: Vec<MemoryEvent>,
pub hw_counters: HardwareCounters,
pub metadata: HashMap<String, String>,
pub status: SessionStatus,
}
#[derive(Debug, Clone, PartialEq)]
pub enum SessionStatus {
Active,
Completed,
Failed(String),
Cancelled,
}
#[derive(Debug, Clone)]
pub enum PerformanceEvent {
FunctionEntry {
function_name: String,
timestamp: Instant,
thread_id: u64,
address: u64,
},
FunctionExit {
function_name: String,
timestamp: Instant,
thread_id: u64,
duration: Duration,
},
KernelLaunch {
kernel_name: String,
timestamp: Instant,
grid_size: (u32, u32, u32),
block_size: (u32, u32, u32),
},
KernelComplete {
kernel_name: String,
timestamp: Instant,
duration: Duration,
occupancy: f32,
},
MemoryAlloc {
size: usize,
address: u64,
timestamp: Instant,
alignment: usize,
},
MemoryFree { address: u64, timestamp: Instant },
CacheMiss {
level: u8,
address: u64,
timestamp: Instant,
},
BranchMisprediction {
address: u64,
timestamp: Instant,
target_address: u64,
},
Custom {
name: String,
timestamp: Instant,
data: HashMap<String, String>,
},
}
#[derive(Debug, Clone)]
pub struct CallStack {
pub timestamp: Instant,
pub thread_id: u64,
pub frames: Vec<StackFrame>,
pub depth: usize,
}
#[derive(Debug, Clone)]
pub struct StackFrame {
pub function_name: String,
pub address: u64,
pub source_location: Option<SourceLocation>,
pub inlined: bool,
pub module_name: String,
}
#[derive(Debug, Clone)]
pub struct SourceLocation {
pub file: String,
pub line: u32,
pub column: u32,
}
#[derive(Debug, Clone)]
pub enum MemoryEvent {
Alloc {
size: usize,
address: u64,
timestamp: Instant,
stack_trace: Vec<StackFrame>,
},
Free {
address: u64,
timestamp: Instant,
stack_trace: Vec<StackFrame>,
},
Access {
address: u64,
size: usize,
is_write: bool,
timestamp: Instant,
},
PageFault {
address: u64,
timestamp: Instant,
fault_type: PageFaultType,
},
}
#[derive(Debug, Clone)]
pub enum PageFaultType {
Major,
Minor,
Protection,
}
#[derive(Debug, Clone, Default)]
pub struct HardwareCounters {
pub cycles: u64,
pub instructions: u64,
pub cache_misses: [u64; 3],
pub cache_references: u64,
pub branch_mispredictions: u64,
pub branches: u64,
pub page_faults: u64,
pub context_switches: u64,
pub cpu_migrations: u64,
pub custom_counters: HashMap<String, u64>,
}
#[derive(Debug, Clone, Default)]
pub struct PerformanceCounters {
pub total_compile_time: Duration,
pub total_execution_time: Duration,
pub compilation_count: u64,
pub execution_count: u64,
pub memory_stats: MemoryStats,
pub function_calls: HashMap<String, u64>,
pub kernel_launches: HashMap<String, u64>,
pub error_counts: HashMap<String, u64>,
}
#[derive(Debug, Clone, Default)]
pub struct MemoryStats {
pub current_usage: usize,
pub peak_usage: usize,
pub total_allocations: u64,
pub total_deallocations: u64,
pub total_bytes_allocated: u64,
pub total_bytes_freed: u64,
pub avg_allocation_size: f64,
pub allocation_histogram: HashMap<usize, u64>,
}
#[derive(Debug, Clone)]
pub struct ProfilerConfig {
pub enabled: bool,
pub sampling_frequency: u32,
pub collect_call_stacks: bool,
pub track_memory: bool,
pub collect_hardware_counters: bool,
pub max_events_per_session: usize,
pub enable_external_profilers: bool,
pub output_format: ProfilerOutputFormat,
pub output_directory: String,
}
#[derive(Debug, Clone)]
pub enum ProfilerOutputFormat {
ChromeTracing,
PerfData,
VTune,
Json,
Binary,
}
#[derive(Debug)]
pub struct SamplingProfiler {
thread_handle: Option<std::thread::JoinHandle<()>>,
config: SamplingConfig,
samples: Arc<Mutex<Vec<Sample>>>,
running: Arc<Mutex<bool>>,
}
#[derive(Debug, Clone)]
pub struct SamplingConfig {
pub interval: Duration,
pub collect_stacks: bool,
pub max_stack_depth: usize,
pub target_threads: Vec<u64>,
}
#[derive(Debug, Clone)]
pub struct Sample {
pub timestamp: Instant,
pub thread_id: u64,
pub cpu_id: u32,
pub pc: u64,
pub stack_trace: Option<Vec<StackFrame>>,
pub cpu_utilization: f32,
pub memory_usage: usize,
}
pub trait ExternalProfiler: Send + Sync + std::fmt::Debug {
fn start(&mut self) -> JitResult<()>;
fn stop(&mut self) -> JitResult<()>;
fn add_function(&mut self, name: &str, address: u64, size: usize) -> JitResult<()>;
fn remove_function(&mut self, address: u64) -> JitResult<()>;
fn export_data(&self, output_path: &str) -> JitResult<()>;
fn name(&self) -> &str;
}
#[derive(Debug)]
pub struct PerfProfiler {
active: bool,
function_map: HashMap<u64, String>,
jit_dump_file: Option<std::fs::File>,
map_file: Option<std::path::PathBuf>,
}
#[derive(Debug)]
pub struct VTuneProfiler {
active: bool,
function_map: HashMap<u64, String>,
}
#[derive(Debug, Clone, Default)]
pub struct ProfilerStats {
pub total_sessions: u64,
pub active_sessions: u64,
pub total_events: u64,
pub total_samples: u64,
pub overhead_percentage: f32,
pub export_count: u64,
}
impl Default for ProfilerConfig {
fn default() -> Self {
Self {
enabled: true,
sampling_frequency: 1000, collect_call_stacks: true,
track_memory: true,
collect_hardware_counters: false, max_events_per_session: 1_000_000,
enable_external_profilers: false,
output_format: ProfilerOutputFormat::Json,
output_directory: std::env::temp_dir()
.join("torsh_profiling")
.display()
.to_string(),
}
}
}
impl ProfilerManager {
pub fn new(config: ProfilerConfig) -> Self {
Self {
sessions: Arc::new(Mutex::new(IndexMap::new())),
counters: Arc::new(Mutex::new(PerformanceCounters::default())),
config,
external_profilers: Vec::new(),
sampling_profiler: None,
stats: Arc::new(Mutex::new(ProfilerStats::default())),
session_counter: Arc::new(Mutex::new(0)),
}
}
pub fn with_defaults() -> Self {
Self::new(ProfilerConfig::default())
}
pub fn start_session(&mut self, name: &str) -> JitResult<String> {
if !self.config.enabled {
return Err(JitError::RuntimeError("Profiling disabled".to_string()));
}
let session_id = {
let mut counter = self
.session_counter
.lock()
.expect("lock should not be poisoned");
*counter += 1;
format!("session_{}", *counter)
};
let session = ProfilingSession {
id: session_id.clone(),
name: name.to_string(),
start_time: Instant::now(),
duration: None,
events: Vec::new(),
call_stacks: Vec::new(),
memory_events: Vec::new(),
hw_counters: HardwareCounters::default(),
metadata: HashMap::new(),
status: SessionStatus::Active,
};
{
let mut sessions = self.sessions.lock().expect("lock should not be poisoned");
sessions.insert(session_id.clone(), session);
}
{
let mut stats = self.stats.lock().expect("lock should not be poisoned");
stats.total_sessions += 1;
stats.active_sessions += 1;
}
if self.config.enable_external_profilers {
for profiler in &mut self.external_profilers {
profiler.start()?;
}
}
Ok(session_id)
}
pub fn stop_session(&mut self, session_id: &str) -> JitResult<()> {
let mut sessions = self.sessions.lock().expect("lock should not be poisoned");
if let Some(session) = sessions.get_mut(session_id) {
session.duration = Some(session.start_time.elapsed());
session.status = SessionStatus::Completed;
let mut stats = self.stats.lock().expect("lock should not be poisoned");
stats.active_sessions = stats.active_sessions.saturating_sub(1);
} else {
return Err(JitError::RuntimeError(format!(
"Session {} not found",
session_id
)));
}
let active_count = {
let stats = self.stats.lock().expect("lock should not be poisoned");
stats.active_sessions
};
if active_count == 0 && self.config.enable_external_profilers {
for profiler in &mut self.external_profilers {
profiler.stop()?;
}
}
Ok(())
}
pub fn record_event(&mut self, session_id: &str, event: PerformanceEvent) -> JitResult<()> {
let mut sessions = self.sessions.lock().expect("lock should not be poisoned");
if let Some(session) = sessions.get_mut(session_id) {
if session.events.len() < self.config.max_events_per_session {
session.events.push(event);
let mut stats = self.stats.lock().expect("lock should not be poisoned");
stats.total_events += 1;
}
}
Ok(())
}
pub fn record_call_stack(&mut self, session_id: &str, call_stack: CallStack) -> JitResult<()> {
if !self.config.collect_call_stacks {
return Ok(());
}
let mut sessions = self.sessions.lock().expect("lock should not be poisoned");
if let Some(session) = sessions.get_mut(session_id) {
session.call_stacks.push(call_stack);
}
Ok(())
}
pub fn start_sampling(&mut self) -> JitResult<()> {
if self.sampling_profiler.is_some() {
return Err(JitError::RuntimeError(
"Sampling profiler already running".to_string(),
));
}
let config = SamplingConfig {
interval: Duration::from_nanos(1_000_000_000 / self.config.sampling_frequency as u64),
collect_stacks: self.config.collect_call_stacks,
max_stack_depth: 64,
target_threads: Vec::new(),
};
let mut sampling_profiler = SamplingProfiler::new(config)?;
sampling_profiler.start()?;
self.sampling_profiler = Some(sampling_profiler);
Ok(())
}
pub fn stop_sampling(&mut self) -> JitResult<()> {
if let Some(mut profiler) = self.sampling_profiler.take() {
profiler.stop()?;
}
Ok(())
}
pub fn add_external_profiler(&mut self, profiler: Box<dyn ExternalProfiler>) {
self.external_profilers.push(profiler);
}
pub fn export_session_data(&self, session_id: &str, output_path: &str) -> JitResult<()> {
let sessions = self.sessions.lock().expect("lock should not be poisoned");
if let Some(session) = sessions.get(session_id) {
match self.config.output_format {
ProfilerOutputFormat::Json => self.export_json(session, output_path)?,
ProfilerOutputFormat::ChromeTracing => {
self.export_chrome_tracing(session, output_path)?
}
_ => {
return Err(JitError::RuntimeError(
"Unsupported export format".to_string(),
))
}
}
let mut stats = self.stats.lock().expect("lock should not be poisoned");
stats.export_count += 1;
}
Ok(())
}
fn export_json(&self, session: &ProfilingSession, output_path: &str) -> JitResult<()> {
use serde_json::json;
let events_json: Vec<_> = session
.events
.iter()
.map(|event| {
json!({
"event": format!("{:?}", event)
})
})
.collect();
let session_json = json!({
"name": session.name,
"id": session.id,
"start_time": session.start_time.elapsed().as_micros(),
"duration": session.duration.map(|d| d.as_micros()),
"events": events_json,
"event_count": session.events.len(),
"call_stack_count": session.call_stacks.len(),
"memory_event_count": session.memory_events.len()
});
std::fs::write(
output_path,
serde_json::to_string_pretty(&session_json)
.map_err(|e| JitError::RuntimeError(format!("JSON serialization failed: {}", e)))?,
)
.map_err(|e| JitError::RuntimeError(format!("Failed to write JSON: {}", e)))?;
Ok(())
}
fn export_chrome_tracing(
&self,
session: &ProfilingSession,
output_path: &str,
) -> JitResult<()> {
use serde_json::json;
let mut trace_events = Vec::new();
for event in &session.events {
match event {
PerformanceEvent::FunctionEntry {
function_name,
timestamp,
thread_id,
..
} => {
trace_events.push(json!({
"name": function_name,
"cat": "function",
"ph": "B", "ts": timestamp.elapsed().as_micros() as u64,
"pid": 1,
"tid": thread_id
}));
}
PerformanceEvent::FunctionExit {
function_name,
timestamp,
thread_id,
duration,
} => {
trace_events.push(json!({
"name": function_name,
"cat": "function",
"ph": "E", "ts": timestamp.elapsed().as_micros() as u64,
"pid": 1,
"tid": thread_id,
"args": { "duration_us": duration.as_micros() }
}));
}
PerformanceEvent::MemoryAlloc {
size,
timestamp,
address,
..
} => {
trace_events.push(json!({
"name": "Memory Allocation",
"cat": "memory",
"ph": "i", "ts": timestamp.elapsed().as_micros() as u64,
"pid": 1,
"tid": 1,
"s": "g", "args": { "size": size, "address": format!("0x{:x}", address) }
}));
}
PerformanceEvent::CacheMiss {
level, timestamp, ..
} => {
trace_events.push(json!({
"name": format!("L{} Cache Miss", level),
"cat": "cache",
"ph": "i",
"ts": timestamp.elapsed().as_micros() as u64,
"pid": 1,
"tid": 1,
"s": "t", }));
}
_ => {
trace_events.push(json!({
"name": format!("{:?}", event),
"cat": "general",
"ph": "i",
"ts": 0,
"pid": 1,
"tid": 1
}));
}
}
}
let metadata = json!({
"process_name": { "1": session.name.clone() },
"thread_name": { "1": "Main Thread" }
});
let trace = json!({
"traceEvents": trace_events,
"displayTimeUnit": "ms",
"metadata": metadata
});
std::fs::write(
output_path,
serde_json::to_string_pretty(&trace)
.map_err(|e| JitError::RuntimeError(format!("JSON serialization failed: {}", e)))?,
)
.map_err(|e| JitError::RuntimeError(format!("Failed to write tracing data: {}", e)))?;
Ok(())
}
pub fn get_session(&self, session_id: &str) -> Option<ProfilingSession> {
let sessions = self.sessions.lock().expect("lock should not be poisoned");
sessions.get(session_id).cloned()
}
pub fn get_counters(&self) -> PerformanceCounters {
let counters = self.counters.lock().expect("lock should not be poisoned");
counters.clone()
}
pub fn get_stats(&self) -> ProfilerStats {
let stats = self.stats.lock().expect("lock should not be poisoned");
stats.clone()
}
pub fn update_counters<F>(&self, update_fn: F)
where
F: FnOnce(&mut PerformanceCounters),
{
let mut counters = self.counters.lock().expect("lock should not be poisoned");
update_fn(&mut *counters);
}
}
impl SamplingProfiler {
pub fn new(config: SamplingConfig) -> JitResult<Self> {
Ok(Self {
thread_handle: None,
config,
samples: Arc::new(Mutex::new(Vec::new())),
running: Arc::new(Mutex::new(false)),
})
}
pub fn start(&mut self) -> JitResult<()> {
let running = self.running.clone();
let samples = self.samples.clone();
let config = self.config.clone();
*running.lock().expect("lock should not be poisoned") = true;
let thread_handle = std::thread::spawn(move || {
let mut last_sample_time = Instant::now();
while *running.lock().expect("lock should not be poisoned") {
let now = Instant::now();
let sample = Sample {
timestamp: now,
thread_id: Self::get_thread_id(),
cpu_id: Self::get_cpu_id(),
pc: 0, stack_trace: Self::collect_stack_trace(),
cpu_utilization: Self::calculate_cpu_utilization(&last_sample_time, &now),
memory_usage: Self::get_memory_usage(),
};
{
let mut samples_guard = samples.lock().expect("lock should not be poisoned");
samples_guard.push(sample);
}
last_sample_time = now;
std::thread::sleep(config.interval);
}
});
self.thread_handle = Some(thread_handle);
Ok(())
}
pub fn stop(&mut self) -> JitResult<()> {
*self.running.lock().expect("lock should not be poisoned") = false;
if let Some(handle) = self.thread_handle.take() {
handle.join().map_err(|_| {
JitError::RuntimeError("Failed to join sampling thread".to_string())
})?;
}
Ok(())
}
pub fn get_samples(&self) -> Vec<Sample> {
let samples = self.samples.lock().expect("lock should not be poisoned");
samples.clone()
}
fn get_thread_id() -> u64 {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let mut hasher = DefaultHasher::new();
std::thread::current().id().hash(&mut hasher);
hasher.finish()
}
fn get_cpu_id() -> u32 {
0
}
fn collect_stack_trace() -> Option<Vec<StackFrame>> {
Some(vec![StackFrame {
function_name: "sampling_thread".to_string(),
address: 0,
source_location: Some(SourceLocation {
file: "profiler.rs".to_string(),
line: 0,
column: 0,
}),
inlined: false,
module_name: "torsh_jit".to_string(),
}])
}
fn calculate_cpu_utilization(_last_time: &Instant, _current_time: &Instant) -> f32 {
(num_cpus::get() as f32) * 0.5 }
fn get_memory_usage() -> usize {
#[cfg(target_os = "linux")]
{
if let Ok(status) = std::fs::read_to_string("/proc/self/status") {
for line in status.lines() {
if line.starts_with("VmRSS:") {
if let Some(kb_str) = line.split_whitespace().nth(1) {
if let Ok(kb) = kb_str.parse::<usize>() {
return kb * 1024; }
}
}
}
}
}
#[cfg(target_os = "macos")]
{
use std::process::Command;
if let Ok(output) = Command::new("ps")
.args(["-o", "rss=", "-p", &std::process::id().to_string()])
.output()
{
if let Ok(text) = String::from_utf8(output.stdout) {
if let Ok(kb) = text.trim().parse::<usize>() {
return kb * 1024; }
}
}
}
0 }
}
impl ExternalProfiler for PerfProfiler {
fn start(&mut self) -> JitResult<()> {
#[cfg(target_os = "linux")]
{
let pid = std::process::id();
let map_file = format!("/tmp/perf-{}.map", pid);
std::fs::write(&map_file, "")
.map_err(|e| JitError::RuntimeError(format!("Failed to create perf map: {}", e)))?;
self.map_file = Some(map_file.into());
}
self.active = true;
Ok(())
}
fn stop(&mut self) -> JitResult<()> {
self.active = false;
Ok(())
}
fn add_function(&mut self, name: &str, address: u64, size: usize) -> JitResult<()> {
self.function_map.insert(address, name.to_string());
#[cfg(target_os = "linux")]
{
if let Some(ref map_file) = self.map_file {
use std::io::Write;
let mut file = std::fs::OpenOptions::new()
.append(true)
.open(map_file)
.map_err(|e| {
JitError::RuntimeError(format!("Failed to open perf map: {}", e))
})?;
writeln!(file, "{:x} {:x} {}", address, size, name).map_err(|e| {
JitError::RuntimeError(format!("Failed to write to perf map: {}", e))
})?;
}
}
Ok(())
}
fn remove_function(&mut self, address: u64) -> JitResult<()> {
self.function_map.remove(&address);
Ok(())
}
fn export_data(&self, output_path: &str) -> JitResult<()> {
use std::io::Write;
let mut file = std::fs::File::create(output_path)
.map_err(|e| JitError::RuntimeError(format!("Failed to create export file: {}", e)))?;
for (address, name) in &self.function_map {
writeln!(file, "{:x} 0 {}", address, name)
.map_err(|e| JitError::RuntimeError(format!("Failed to write perf data: {}", e)))?;
}
Ok(())
}
fn name(&self) -> &str {
"PerfProfiler"
}
}
impl ExternalProfiler for VTuneProfiler {
fn start(&mut self) -> JitResult<()> {
#[cfg(target_os = "linux")]
{
log::info!("VTune profiler started (stub implementation)");
}
#[cfg(target_os = "windows")]
{
log::info!("VTune profiler started on Windows (stub implementation)");
}
self.active = true;
Ok(())
}
fn stop(&mut self) -> JitResult<()> {
self.active = false;
Ok(())
}
fn add_function(&mut self, name: &str, address: u64, size: usize) -> JitResult<()> {
self.function_map.insert(address, name.to_string());
if self.active {
log::debug!(
"Registered function '{}' at 0x{:x} (size: {}) with VTune",
name,
address,
size
);
}
Ok(())
}
fn remove_function(&mut self, address: u64) -> JitResult<()> {
self.function_map.remove(&address);
if self.active {
log::debug!("Unregistered function at 0x{:x} from VTune", address);
}
Ok(())
}
fn export_data(&self, output_path: &str) -> JitResult<()> {
use std::io::Write;
let mut file = std::fs::File::create(output_path)
.map_err(|e| JitError::RuntimeError(format!("Failed to create export file: {}", e)))?;
writeln!(file, "# VTune JIT Symbol Map")
.map_err(|e| JitError::RuntimeError(format!("Write failed: {}", e)))?;
writeln!(file, "# Format: <address> <size> <name>")
.map_err(|e| JitError::RuntimeError(format!("Write failed: {}", e)))?;
writeln!(file).map_err(|e| JitError::RuntimeError(format!("Write failed: {}", e)))?;
for (address, name) in &self.function_map {
writeln!(file, "{:016x} 0000 {}", address, name).map_err(|e| {
JitError::RuntimeError(format!("Failed to write VTune data: {}", e))
})?;
}
Ok(())
}
fn name(&self) -> &str {
"VTuneProfiler"
}
}
impl PerfProfiler {
pub fn new() -> Self {
Self {
active: false,
function_map: HashMap::new(),
jit_dump_file: None,
map_file: None,
}
}
}
impl VTuneProfiler {
pub fn new() -> Self {
Self {
active: false,
function_map: HashMap::new(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_profiler_manager_creation() {
let manager = ProfilerManager::with_defaults();
assert!(manager.config.enabled);
assert_eq!(manager.config.sampling_frequency, 1000);
}
#[test]
fn test_session_lifecycle() {
let mut manager = ProfilerManager::with_defaults();
let session_id = manager.start_session("test_session").unwrap();
assert!(!session_id.is_empty());
let session = manager.get_session(&session_id).unwrap();
assert_eq!(session.name, "test_session");
assert_eq!(session.status, SessionStatus::Active);
manager.stop_session(&session_id).unwrap();
let session = manager.get_session(&session_id).unwrap();
assert_eq!(session.status, SessionStatus::Completed);
assert!(session.duration.is_some());
}
#[test]
fn test_performance_event_recording() {
let mut manager = ProfilerManager::with_defaults();
let session_id = manager.start_session("test_session").unwrap();
let event = PerformanceEvent::FunctionEntry {
function_name: "test_function".to_string(),
timestamp: Instant::now(),
thread_id: 1,
address: 0x1000,
};
manager.record_event(&session_id, event).unwrap();
let session = manager.get_session(&session_id).unwrap();
assert_eq!(session.events.len(), 1);
}
#[test]
fn test_external_profiler_integration() {
let mut manager = ProfilerManager::with_defaults();
let perf_profiler = Box::new(PerfProfiler::new());
manager.add_external_profiler(perf_profiler);
assert_eq!(manager.external_profilers.len(), 1);
}
}