puffin_http server: no need to call update() each frame

Author: emilk

puffin/src/lib.rs

@@ -512,6 +512,12 @@ impl Ord for OrderedData {
     }
 }
 
+/// Add these to [`GlobalProfiler`] with [`GlobalProfiler::add_sink`].
+pub type FrameSink = Box<dyn Fn(Arc<FrameData>) + Send>;
+
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
+pub struct FrameSinkId(u64);
+
 /// Singleton. Collects profiling data from multiple threads.
 pub struct GlobalProfiler {
     current_frame_index: FrameIndex,
@@ -523,6 +529,9 @@ pub struct GlobalProfiler {
 
     slowest_frames: std::collections::BinaryHeap<OrderedData>,
     max_slow: usize,
+
+    next_sink_id: FrameSinkId,
+    sinks: std::collections::HashMap<FrameSinkId, FrameSink>,
 }
 
 impl Default for GlobalProfiler {
@@ -537,6 +546,8 @@ impl Default for GlobalProfiler {
             max_recent,
             slowest_frames: std::collections::BinaryHeap::with_capacity(max_slow),
             max_slow,
+            next_sink_id: FrameSinkId(1),
+            sinks: Default::default(),
         }
     }
 }
@@ -573,6 +584,10 @@ impl GlobalProfiler {
 
     /// Manually add frame data.
     pub fn add_frame(&mut self, new_frame: Arc<FrameData>) {
+        for sink in self.sinks.values() {
+            sink(new_frame.clone());
+        }
+
         let add_to_slowest = if self.slowest_frames.len() < self.max_slow {
             true
         } else if let Some(fastest_of_the_slow) = self.slowest_frames.peek() {
@@ -703,6 +718,19 @@ impl GlobalProfiler {
 
         Ok(slf)
     }
+
+    /// Call this function with each new finished frame.
+    /// The returned [`FrameSinkId`] can be used to remove the sink with [`Self::remove_sink`].
+    pub fn add_sink(&mut self, sink: FrameSink) -> FrameSinkId {
+        let id = self.next_sink_id;
+        self.next_sink_id.0 += 1;
+        self.sinks.insert(id, sink);
+        id
+    }
+
+    pub fn remove_sink(&mut self, id: FrameSinkId) -> Option<FrameSink> {
+        self.sinks.remove(&id)
+    }
 }
 
 // ----------------------------------------------------------------------------

puffin_http/examples/server.rs

@@ -7,7 +7,7 @@ fn main() {
     let server_addr = format!("0.0.0.0:{}", puffin_http::DEFAULT_PORT);
     eprintln!("Serving demo profile data on {}", server_addr);
 
-    let puffin_server = puffin_http::Server::new(&server_addr).unwrap();
+    let _puffin_server = puffin_http::Server::new(&server_addr).unwrap();
 
     puffin::set_scopes_on(true);
 
@@ -16,7 +16,6 @@ fn main() {
     loop {
         puffin::profile_scope!("main_loop");
         puffin::GlobalProfiler::lock().new_frame();
-        puffin_server.update();
 
         // Give us something to inspect:
 

puffin_http/src/server.rs

@@ -1,4 +1,5 @@
 use anyhow::Context as _;
+use puffin::GlobalProfiler;
 use std::{
     io::Write,
     net::{SocketAddr, TcpListener, TcpStream},
@@ -7,8 +8,10 @@ use std::{
 
 /// Listens for incoming connections
 /// and streams them puffin profiler data.
+///
+/// Drop to stop transmitting and listening for new connections.
 pub struct Server {
-    tx: std::sync::mpsc::Sender<Arc<puffin::FrameData>>,
+    sink_id: puffin::FrameSinkId,
 }
 
 impl Server {
@@ -19,9 +22,8 @@ impl Server {
             .set_nonblocking(true)
             .context("TCP set_nonblocking")?;
 
-        let (tx, rx) = std::sync::mpsc::channel();
-
-        let server = Server { tx };
+        let (tx, rx): (std::sync::mpsc::Sender<Arc<puffin::FrameData>>, _) =
+            std::sync::mpsc::channel();
 
         std::thread::Builder::new()
             .name("puffin-server".to_owned())
@@ -42,15 +44,17 @@ impl Server {
             })
             .context("Couldn't spawn thread")?;
 
-        Ok(server)
+        let sink_id = GlobalProfiler::lock().add_sink(Box::new(move |frame| {
+            tx.send(frame).ok();
+        }));
+
+        Ok(Server { sink_id })
     }
+}
 
-    /// Call this once per frame, right after calling [`puffin::GlobalProfiler::new_frame`].
-    pub fn update(&self) {
-        let latest_frame = puffin::GlobalProfiler::lock().latest_frame();
-        if let Some(latest_frame) = latest_frame {
-            self.tx.send(latest_frame).ok();
-        }
+impl Drop for Server {
+    fn drop(&mut self) {
+        GlobalProfiler::lock().remove_sink(self.sink_id);
     }
 }