rrtrace 0.1.0

data/sig/rrtrace.rbs

@@ -0,0 +1,4 @@
+module Rrtrace
+  VERSION: String
+  def self.visualizer_path: () -> String
+end

data/src/main.rs

@@ -0,0 +1,197 @@
+use crate::renderer::Renderer;
+use crate::ringbuffer::{EventRingBuffer, RRTraceEvent};
+use crate::trace_state::{FastTrace, SlowTrace, VISIBLE_DURATION};
+use std::ffi::CString;
+use std::sync::atomic::AtomicU64;
+use std::sync::{Arc, atomic};
+use std::{env, mem, thread};
+use winit::application::ApplicationHandler;
+use winit::event::*;
+use winit::event_loop::{ControlFlow, EventLoop};
+use winit::window::Window;
+
+mod renderer;
+mod ringbuffer;
+#[cfg_attr(unix, path = "shm_unix.rs")]
+#[cfg_attr(windows, path = "shm_windows.rs")]
+mod shm;
+mod trace_state;
+
+struct App {
+    window: Option<Arc<Window>>,
+    renderer: Renderer,
+}
+
+impl App {
+    fn new(renderer: Renderer) -> Self {
+        Self {
+            window: None,
+            renderer,
+        }
+    }
+}
+
+impl ApplicationHandler for App {
+    fn resumed(&mut self, event_loop: &winit::event_loop::ActiveEventLoop) {
+        let window = Arc::new(
+            event_loop
+                .create_window(Window::default_attributes().with_title("rrtrace visualizer"))
+                .unwrap(),
+        );
+        self.renderer.set_window(window.clone());
+        self.window = Some(window);
+    }
+
+    fn window_event(
+        &mut self,
+        event_loop: &winit::event_loop::ActiveEventLoop,
+        _window_id: winit::window::WindowId,
+        event: WindowEvent,
+    ) {
+        let Some(window) = self.window.as_ref() else {
+            return;
+        };
+
+        match event {
+            WindowEvent::CloseRequested
+            | WindowEvent::KeyboardInput {
+                event:
+                    KeyEvent {
+                        state: ElementState::Pressed,
+                        logical_key: winit::keyboard::Key::Named(winit::keyboard::NamedKey::Escape),
+                        ..
+                    },
+                ..
+            } => event_loop.exit(),
+            WindowEvent::Resized(physical_size) => {
+                self.renderer.resize(physical_size);
+            }
+            WindowEvent::RedrawRequested => match self.renderer.render() {
+                Ok(_) => {}
+                Err(wgpu::SurfaceError::Lost) => self.renderer.resize(window.inner_size()),
+                Err(wgpu::SurfaceError::OutOfMemory) => event_loop.exit(),
+                Err(e) => eprintln!("{:?}", e),
+            },
+            _ => {}
+        }
+    }
+
+    fn about_to_wait(&mut self, _event_loop: &winit::event_loop::ActiveEventLoop) {
+        let updated = self.renderer.sync();
+        if updated && let Some(window) = self.window.as_ref() {
+            window.request_redraw();
+        }
+    }
+}
+
+fn main() {
+    assert_eq!(env::args().len(), 2, "Usage: rrtrace <shm_name>");
+    let shm_name = env::args().nth(1).unwrap();
+
+    let (instance, adapter, device, queue) = pollster::block_on(init_gpu());
+    let event_queue = Arc::new(crossbeam_queue::SegQueue::new());
+    let result_queue = Arc::new(crossbeam_queue::SegQueue::new());
+    thread::Builder::new()
+        .name("queue pipe".to_owned())
+        .spawn(queue_pipe_thread(shm_name, Arc::clone(&event_queue)))
+        .unwrap();
+    thread::Builder::new()
+        .name("trace".to_owned())
+        .spawn(trace_thread(
+            Arc::clone(&event_queue),
+            Arc::clone(&result_queue),
+        ))
+        .unwrap();
+
+    let event_loop = EventLoop::new().unwrap();
+    event_loop.set_control_flow(ControlFlow::Poll);
+    let mut app = App::new(Renderer::new(
+        instance,
+        adapter,
+        device,
+        queue,
+        result_queue,
+    ));
+    event_loop.run_app(&mut app).unwrap();
+}
+
+async fn init_gpu() -> (wgpu::Instance, wgpu::Adapter, wgpu::Device, wgpu::Queue) {
+    let instance = wgpu::Instance::default();
+    let adapter = instance
+        .request_adapter(&wgpu::RequestAdapterOptions {
+            power_preference: wgpu::PowerPreference::HighPerformance,
+            compatible_surface: None,
+            force_fallback_adapter: false,
+        })
+        .await
+        .unwrap();
+    let (device, queue) = adapter
+        .request_device(&wgpu::DeviceDescriptor {
+            label: None,
+            required_features: wgpu::Features::empty(),
+            required_limits: wgpu::Limits::default(),
+            experimental_features: Default::default(),
+            memory_hints: Default::default(),
+            trace: Default::default(),
+        })
+        .await
+        .unwrap();
+    (instance, adapter, device, queue)
+}
+
+fn queue_pipe_thread(
+    shm_name: String,
+    event_queue: Arc<crossbeam_queue::SegQueue<Vec<RRTraceEvent>>>,
+) -> impl FnOnce() + Send + 'static {
+    move || {
+        let shm = unsafe {
+            shm::SharedMemory::open(
+                CString::new(shm_name).unwrap(),
+                mem::size_of::<ringbuffer::RRTraceEventRingBuffer>(),
+            )
+        };
+        let mut ringbuffer = unsafe { EventRingBuffer::new(shm.as_ptr(), move || drop(shm)) };
+        let mut buffer = vec![Default::default(); 65536];
+        loop {
+            let count = ringbuffer.read(&mut buffer);
+            if count > 0 {
+                buffer.truncate(count);
+                event_queue.push(buffer.clone());
+                buffer.resize_with(65536, Default::default);
+            }
+        }
+    }
+}
+fn trace_thread(
+    event_queue: Arc<crossbeam_queue::SegQueue<Vec<RRTraceEvent>>>,
+    result_queue: Arc<crossbeam_queue::SegQueue<SlowTrace>>,
+) -> impl FnOnce() + Send + 'static {
+    move || {
+        static LATEST_END_TIME: AtomicU64 = AtomicU64::new(0);
+        let mut start_time = 0u64;
+        let mut fast_trace = FastTrace::new();
+        loop {
+            let Some(events) = event_queue.pop() else {
+                continue;
+            };
+            rayon_core::spawn({
+                let fast_trace = fast_trace.clone();
+                let events = events.clone();
+                let result_queue = result_queue.clone();
+                move || {
+                    if start_time + VISIBLE_DURATION
+                        < LATEST_END_TIME.load(atomic::Ordering::Relaxed)
+                    {
+                        return;
+                    }
+                    let slow_trace = SlowTrace::trace(start_time, fast_trace, &events);
+                    result_queue.push(slow_trace);
+                }
+            });
+            fast_trace.process_events(&events);
+            let end_time = events.last().unwrap().timestamp();
+            LATEST_END_TIME.store(end_time, atomic::Ordering::Relaxed);
+            start_time = end_time;
+        }
+    }
+}

data/src/renderer/vertex_arena.rs

@@ -0,0 +1,305 @@
+use bytemuck::NoUninit;
+use std::collections::{BTreeMap, BTreeSet, HashMap};
+use std::fmt;
+use std::fmt::{Debug, Formatter};
+use std::ops::Range;
+use std::sync::atomic;
+use wgpu::{Buffer, BufferAddress, BufferDescriptor, BufferUsages, Device, Queue};
+
+#[derive(Debug, Clone, Copy, Ord, PartialOrd, Eq, PartialEq, Hash)]
+pub struct AllocationId(usize);
+
+impl AllocationId {
+    fn new() -> AllocationId {
+        static COUNTER: atomic::AtomicUsize = atomic::AtomicUsize::new(0);
+        AllocationId(COUNTER.fetch_add(1, atomic::Ordering::Relaxed))
+    }
+}
+
+pub struct VertexArena<T> {
+    device: Device,
+    queue: Queue,
+    data: Vec<T>,
+    gpu_buffer: Vec<Buffer>,
+    max_buffer_size: u64,
+    allocations: HashMap<AllocationId, Range<usize>>,
+    free_list: FreeList,
+    dirty_range: Range<usize>,
+}
+
+struct FreeList {
+    by_start: BTreeMap<usize, usize>,
+    by_size: BTreeSet<(usize, usize)>,
+}
+
+impl FreeList {
+    fn new() -> Self {
+        Self {
+            by_start: BTreeMap::new(),
+            by_size: BTreeSet::new(),
+        }
+    }
+
+    fn alloc(&mut self, len: usize) -> Option<Range<usize>> {
+        let &(size, start) = self.by_size.range((len, 0)..).next()?;
+        self.by_size.remove(&(size, start));
+        self.by_start.remove(&start);
+        if size > len {
+            let new_start = start + len;
+            let new_size = size - len;
+            self.by_start.insert(new_start, new_start + new_size);
+            self.by_size.insert((new_size, new_start));
+        }
+        Some(start..start + len)
+    }
+
+    fn dealloc(&mut self, range: Range<usize>) {
+        let mut start = range.start;
+        let mut end = range.end;
+
+        if let Some((&next_start, &next_end)) = self.by_start.range(end..).next() {
+            if next_start == end {
+                self.by_size.remove(&(next_end - next_start, next_start));
+                self.by_start.remove(&next_start);
+                end = next_end;
+            }
+        }
+
+        if let Some((&prev_start, &prev_end)) = self.by_start.range(..start).next_back() {
+            if prev_end == start {
+                self.by_size.remove(&(prev_end - prev_start, prev_start));
+                self.by_start.remove(&prev_start);
+                start = prev_start;
+            }
+        }
+
+        self.by_start.insert(start, end);
+        self.by_size.insert((end - start, start));
+    }
+}
+
+impl<T> Debug for VertexArena<T>
+where
+    T: Debug,
+{
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.debug_list().entries(&self.data).finish()
+    }
+}
+
+impl<T> VertexArena<T> {
+    pub fn new(device: Device, queue: Queue, usage: BufferUsages) -> VertexArena<T> {
+        let max_buffer_size = device.limits().max_buffer_size;
+        let gpu_buffer = device.create_buffer(&BufferDescriptor {
+            label: None,
+            size: (max_buffer_size / size_of::<T>() as u64).min(256) * size_of::<T>() as u64,
+            usage,
+            mapped_at_creation: false,
+        });
+        VertexArena {
+            data: Vec::new(),
+            device,
+            queue,
+            gpu_buffer: vec![gpu_buffer],
+            max_buffer_size,
+            allocations: HashMap::new(),
+            free_list: FreeList::new(),
+            dirty_range: usize::MAX..0,
+        }
+    }
+
+    pub fn alloc(&mut self, len: usize) -> (AllocationId, &mut [T])
+    where
+        T: Default,
+    {
+        let id = AllocationId::new();
+        let range = if let Some(range) = self.free_list.alloc(len) {
+            range
+        } else {
+            let start = self.data.len();
+            self.data.resize_with(start + len, T::default);
+            start..start + len
+        };
+
+        self.allocations.insert(id, range.clone());
+        self.dirty_range.start = self.dirty_range.start.min(range.start);
+        self.dirty_range.end = self.dirty_range.end.max(range.end);
+
+        let result = &mut self.data[range.clone()];
+        assert_eq!(result.len(), len);
+        (id, result)
+    }
+
+    pub fn dealloc(&mut self, id: AllocationId) {
+        if let Some(range) = self.allocations.remove(&id) {
+            self.free_list.dealloc(range);
+        }
+    }
+
+    pub fn sync(&mut self)
+    where
+        T: NoUninit,
+    {
+        if self.dirty_range.start >= self.dirty_range.end {
+            return;
+        }
+
+        let filled_buffer_len = self.max_buffer_size / size_of::<T>() as u64;
+        let single_buffer_size_max = filled_buffer_len * size_of::<T>() as u64;
+        if let [gpu_buffer] = self.gpu_buffer.as_slice() {
+            let required_size = self.data.len() as u64 * size_of::<T>() as u64;
+            if required_size > gpu_buffer.size() {
+                if required_size <= self.max_buffer_size {
+                    self.gpu_buffer = vec![self.device.create_buffer(&BufferDescriptor {
+                        label: None,
+                        size: required_size.next_power_of_two(),
+                        usage: gpu_buffer.usage(),
+                        mapped_at_creation: false,
+                    })];
+                } else {
+                    let required_buffer_count = required_size.div_ceil(single_buffer_size_max);
+                    let buffer_usages = gpu_buffer.usage();
+                    if gpu_buffer.size() < single_buffer_size_max {
+                        self.gpu_buffer.clear();
+                    }
+                    for _ in 1..required_buffer_count {
+                        self.gpu_buffer
+                            .push(self.device.create_buffer(&BufferDescriptor {
+                                label: None,
+                                size: single_buffer_size_max,
+                                usage: buffer_usages,
+                                mapped_at_creation: false,
+                            }));
+                    }
+                }
+                self.dirty_range = 0..self.data.len();
+            }
+        } else {
+            let new_buffer_len = self.data.len().div_ceil(filled_buffer_len as usize);
+            let usage = self.gpu_buffer[0].usage();
+            for _ in self.gpu_buffer.len()..new_buffer_len {
+                self.gpu_buffer
+                    .push(self.device.create_buffer(&BufferDescriptor {
+                        label: None,
+                        size: single_buffer_size_max,
+                        usage,
+                        mapped_at_creation: false,
+                    }));
+            }
+        }
+
+        if let [gpu_buffer] = self.gpu_buffer.as_slice() {
+            let dirty_data = &self.data[self.dirty_range.clone()];
+            let offset = (self.dirty_range.start * size_of::<T>()) as BufferAddress;
+            let bytes: &[u8] = bytemuck::cast_slice(dirty_data);
+            self.queue.write_buffer(gpu_buffer, offset, bytes);
+        } else {
+            let start_block = self.dirty_range.start / filled_buffer_len as usize;
+            let start_item = self.dirty_range.start % filled_buffer_len as usize;
+            let end_block = self.dirty_range.end / filled_buffer_len as usize;
+            let end_item = self.dirty_range.end % filled_buffer_len as usize;
+            match &self.gpu_buffer[start_block..=end_block] {
+                [] => unreachable!(),
+                [buffer] => {
+                    let dirty_data = &self.data[self.dirty_range.clone()];
+                    let offset = (start_item * size_of::<T>()) as BufferAddress;
+                    let bytes: &[u8] = bytemuck::cast_slice(dirty_data);
+                    self.queue.write_buffer(buffer, offset, bytes);
+                }
+                [first, mid @ .., last] => {
+                    let data = &self.data[start_block * filled_buffer_len as usize
+                        ..((end_block + 1) * filled_buffer_len as usize).min(self.data.len())];
+                    let mut data_iter = data.chunks(filled_buffer_len as usize);
+                    let first_chunk = data_iter.next().unwrap();
+                    let last_chunk = data_iter.next_back().unwrap();
+                    self.queue.write_buffer(
+                        first,
+                        (start_item * size_of::<T>()) as BufferAddress,
+                        bytemuck::cast_slice(&first_chunk[start_item..]),
+                    );
+                    if end_item > 0 {
+                        self.queue.write_buffer(
+                            last,
+                            0,
+                            bytemuck::cast_slice(&last_chunk[..end_item]),
+                        );
+                    }
+                    for (buffer, data) in mid.iter().zip(data_iter) {
+                        self.queue
+                            .write_buffer(buffer, 0, bytemuck::cast_slice(data));
+                    }
+                }
+            }
+        }
+
+        self.dirty_range = usize::MAX..0;
+    }
+
+    pub fn read_buffers(&self, mut f: impl FnMut(&Buffer, usize)) {
+        if let [buffer] = &self.gpu_buffer.as_slice() {
+            f(buffer, self.data.len());
+        } else {
+            let filled_buffer_len = self.max_buffer_size as usize / size_of::<T>();
+            let num_buffers = self.data.len() / filled_buffer_len;
+            let buffer_tail = self.data.len() % filled_buffer_len;
+            for buffer in self.gpu_buffer.iter().take(num_buffers) {
+                f(buffer, filled_buffer_len);
+            }
+            if buffer_tail > 0 {
+                f(&self.gpu_buffer[num_buffers], buffer_tail);
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_free_list_merge() {
+        let mut fl = FreeList::new();
+
+        // [10..20]
+        fl.dealloc(10..20);
+        assert_eq!(fl.by_start.get(&10), Some(&20));
+
+        // [0..5, 10..20]
+        fl.dealloc(0..5);
+        assert_eq!(fl.by_start.len(), 2);
+
+        // [0..5, 10..20, 25..30]
+        fl.dealloc(25..30);
+        assert_eq!(fl.by_start.len(), 3);
+
+        // Merge next: [0..10, 10..20, 25..30] -> [0..20, 25..30]
+        fl.dealloc(5..10);
+        assert_eq!(fl.by_start.len(), 2);
+        assert_eq!(fl.by_start.get(&0), Some(&20));
+
+        // Merge both: [0..20, 20..25, 25..30] -> [0..30]
+        fl.dealloc(20..25);
+        assert_eq!(fl.by_start.len(), 1);
+        assert_eq!(fl.by_start.get(&0), Some(&30));
+    }
+
+    #[test]
+    fn test_free_list_alloc_split() {
+        let mut fl = FreeList::new();
+        fl.dealloc(0..10);
+        fl.dealloc(20..30);
+        fl.dealloc(40..50);
+
+        // Alloc 5. Should pick 0..10
+        let r1 = fl.alloc(5).unwrap();
+        assert_eq!(r1, 0..5);
+        // remains 5..10, 20..30, 40..50
+        assert_eq!(fl.by_start.get(&5), Some(&10));
+
+        // Alloc 10. Should pick 20..30
+        let r2 = fl.alloc(10).unwrap();
+        assert_eq!(r2, 20..30);
+        // remains 5..10, 40..50
+        assert_eq!(fl.by_start.len(), 2);
+    }
+}