Fix puffin integration memory leak

Author: Azkellas

examples/demo.rs

@@ -60,6 +60,8 @@ struct State {
     window: Option<Arc<winit::window::Window>>,
     gfx_state: Option<GfxState>,
     latest_profiler_results: Option<Vec<GpuTimerQueryResult>>,
+    #[cfg(feature = "puffin")]
+    puffin_scope_cache: wgpu_profiler::puffin::PuffinScopeCache,
 }
 
 struct GfxState {
@@ -264,6 +266,7 @@ impl ApplicationHandler<()> for State {
                     let mut gpu_profiler = PUFFIN_GPU_PROFILER.lock().unwrap();
                     wgpu_profiler::puffin::output_frame_to_puffin(
                         &mut gpu_profiler,
+                        &mut self.puffin_scope_cache,
                         self.latest_profiler_results.as_deref().unwrap_or_default(),
                     );
                     gpu_profiler.new_frame();

src/puffin.rs

@@ -1,11 +1,23 @@
-use puffin::{GlobalProfiler, NanoSecond, ScopeDetails, StreamInfo, ThreadInfo};
+use std::collections::HashMap;
+
+use puffin::{GlobalProfiler, NanoSecond, ScopeDetails, ScopeId, StreamInfo, ThreadInfo};
 
 use crate::GpuTimerQueryResult;
 
+/// Cache for registered puffin scope IDs to avoid memory leaks.
+#[derive(Default)]
+pub struct PuffinScopeCache {
+    scope_ids: HashMap<String, ScopeId>,
+}
+
 /// Visualize the query results in a `puffin::GlobalProfiler`.
-pub fn output_frame_to_puffin(profiler: &mut GlobalProfiler, query_result: &[GpuTimerQueryResult]) {
+pub fn output_frame_to_puffin(
+    profiler: &mut GlobalProfiler,
+    cache: &mut PuffinScopeCache,
+    query_result: &[GpuTimerQueryResult],
+) {
     let mut stream_info = StreamInfo::default();
-    collect_stream_info_recursive(profiler, &mut stream_info, query_result, 0);
+    build_stream_info(profiler, cache, &mut stream_info, query_result, 0);
 
     profiler.report_user_scopes(
         ThreadInfo {
@@ -16,29 +28,41 @@ pub fn output_frame_to_puffin(profiler: &mut GlobalProfiler, query_result: &[Gpu
     );
 }
 
-fn collect_stream_info_recursive(
+fn build_stream_info(
     profiler: &mut GlobalProfiler,
+    cache: &mut PuffinScopeCache,
     stream_info: &mut StreamInfo,
     query_result: &[GpuTimerQueryResult],
     depth: usize,
 ) {
-    let details: Vec<_> = query_result
-        .iter()
-        .map(|query| ScopeDetails::from_scope_name(query.label.clone()))
-        .collect();
-    let ids = profiler.register_user_scopes(&details);
-    for (query, id) in query_result.iter().zip(ids) {
+    for query in query_result {
+        // Use get() first to avoid cloning the label on cache hits.
+        let id = if let Some(&id) = cache.scope_ids.get(&query.label) {
+            id
+        } else {
+            let details = [ScopeDetails::from_scope_name(query.label.clone())];
+            let id = profiler.register_user_scopes(&details)[0];
+            cache.scope_ids.insert(query.label.clone(), id);
+            id
+        };
+
         if let Some(time) = &query.time {
             let start = (time.start * 1e9) as NanoSecond;
             let end = (time.end * 1e9) as NanoSecond;
 
             stream_info.depth = stream_info.depth.max(depth);
             stream_info.num_scopes += 1;
             stream_info.range_ns.0 = stream_info.range_ns.0.min(start);
-            stream_info.range_ns.1 = stream_info.range_ns.0.max(end);
+            stream_info.range_ns.1 = stream_info.range_ns.1.max(end);
 
             let (offset, _) = stream_info.stream.begin_scope(|| start, id, "");
-            collect_stream_info_recursive(profiler, stream_info, &query.nested_queries, depth + 1);
+            build_stream_info(
+                profiler,
+                cache,
+                stream_info,
+                &query.nested_queries,
+                depth + 1,
+            );
             stream_info.stream.end_scope(offset, end as NanoSecond);
         }
     }