Alllow using custom GlobalProfiler instance with puffin_http::Server

puffin/src/global_profiler.rs

@@ -127,7 +127,7 @@ impl GlobalProfiler {
     }
 
     /// Reports some profiling data. Called from [`ThreadProfiler`].
-    pub(crate) fn report(
+    pub fn report(
         &mut self,
         info: ThreadInfo,
         scope_details: &[ScopeDetails],

puffin_http/Cargo.toml

@@ -25,3 +25,5 @@ puffin = { version = "0.19.0", path = "../puffin", features = [
 
 [dev-dependencies]
 simple_logger = "4.2"
+paste = "1.0.15"
+once_cell = "1.19.0"

puffin_http/src/server.rs

@@ -21,11 +21,211 @@ pub struct Server {
     sink_id: puffin::FrameSinkId,
     join_handle: Option<std::thread::JoinHandle<()>>,
     num_clients: Arc<AtomicUsize>,
+    sink_remove: fn(puffin::FrameSinkId) -> (),
 }
 
 impl Server {
     /// Start listening for connections on this addr (e.g. "0.0.0.0:8585")
+    ///
+    /// Connects to the [GlobalProfiler]
     pub fn new(bind_addr: &str) -> anyhow::Result<Self> {
+        fn global_add(sink: puffin::FrameSink) -> puffin::FrameSinkId {
+            GlobalProfiler::lock().add_sink(sink)
+        }
+        fn global_remove(id: puffin::FrameSinkId) {
+            GlobalProfiler::lock().remove_sink(id);
+        }
+
+        Self::new_custom(bind_addr, global_add, global_remove)
+    }
+
+    /// Starts a new puffin server, with a custom function for installing the server's sink
+    ///
+    /// # Arguments
+    /// * `bind_addr` - The address to bind to, when listening for connections
+    /// (e.g. "localhost:8585" or "127.0.0.1:8585")
+    /// * `sink_install` - A function that installs the [Server]'s sink into
+    /// a [GlobalProfiler], and then returns the [FrameSinkId] so that the sink can be removed later
+    /// * `sink_remove` - A function that reverts `sink_install`.
+    /// This should be a call to remove the sink from the profiler ([GlobalProfiler::remove_sink])
+    ///
+    /// # Example
+    ///
+    /// Using this is slightly complicated, but it is possible to use this to set a custom profiler per-thread,
+    /// such that threads can be grouped together and profiled separately. E.g. you could have one profiling server
+    /// instance for the main UI loop, and another for the background worker loop, and events/frames from those thread(s)
+    /// would be completely separated. You can then hook up two separate instances of `puffin_viewer` and profile them separately.
+    ///
+    /// ## Per-Thread Profiling
+    /// ```
+    /// # use puffin::GlobalProfiler;
+    /// # use puffin::{StreamInfoRef, ThreadInfo, ScopeDetails};
+    /// # use puffin_http::Server;
+    /// # use puffin::ThreadProfiler;
+    /// #
+    /// # pub fn main() {
+    /// #
+    /// #
+    /// // Initialise the profiling server for the main app
+    /// let default_server = Server::new("localhost:8585").expect("failed to create default profiling server");
+    /// puffin::profile_scope!("main_scope");
+    ///
+    /// // Create a new [GlobalProfiler] instance. This is where we will be sending the events to for our threads.
+    /// // [OnceLock] and [Mutex] are there so that we can safely get exclusive mutable access.
+    /// static CUSTOM_PROFILER: std::sync::OnceLock<std::sync::Mutex<GlobalProfiler>> = std::sync::OnceLock::new();
+    /// // Helper function to access the profiler
+    /// fn get_custom_profiler() -> std::sync::MutexGuard<'static, GlobalProfiler> {
+    ///    CUSTOM_PROFILER.get_or_init(|| std::sync::Mutex::new(GlobalProfiler::default()))
+    ///         .lock().expect("failed to lock custom profiler")
+    /// }
+    /// // Create the custom profiling server that uses our custom profiler instead of the global/default one
+    /// let thread_server = Server::new_custom(
+    ///     "localhost:6969",
+    ///     // Adds the [Server]'s sink to our custom profiler
+    ///     |sink| get_custom_profiler().add_sink(sink),
+    ///     // Remove
+    ///     |id| _ = get_custom_profiler().remove_sink(id)
+    /// );
+    /// 
+    /// // Create some custom threads where we use the custom profiler and server
+    /// std::thread::scope(|scope| {
+    ///     scope.spawn(move ||{
+    ///         // Tell this thread to use the custom profiler
+    ///         let _ = ThreadProfiler::initialize(
+    ///             // Use the same time source as default puffin
+    ///             puffin::now_ns,
+    ///             // However redirect the events to our `custom_profiler`, instead of the default
+    ///             // which would be the one returned by [GlobalProfiler::lock()]
+    ///             |info: ThreadInfo, details: &[ScopeDetails], stream: &StreamInfoRef<'_>|
+    ///                 get_custom_profiler().report(info, details, stream)
+    ///         );
+    ///
+    ///         // Do work
+    ///         {
+    ///             puffin::profile_scope!("inside_thread");
+    ///             println!("hello from the thread");
+    ///             std::thread::sleep(std::time::Duration::from_secs(1));
+    ///         }
+    ///
+    ///         // Tell our profiler that we are done with this frame
+    ///         // This will be sent to the server on port 6969
+    ///         get_custom_profiler().new_frame();
+    ///     });
+    /// });
+    ///
+    /// // New frame for the global profiler. This is completely separate from the scopes with the custom profiler
+    /// GlobalProfiler::lock().new_frame();
+    /// #
+    /// #
+    /// # }
+    /// ```
+    ///
+    /// ## Helpful Macro
+    /// ```rust
+    /// # use std::thread::sleep;
+    /// # use std::time::Duration;
+    /// 
+    /// /// This macro makes it much easier to define profilers
+    /// ///
+    /// /// This macro makes use of the `paste` crate to generate unique identifiers, and `tracing` to log events
+    /// macro_rules! profiler {
+    ///     ($(
+    ///          {name: $name:ident, port: $port:expr $(,install: |$install_var:ident| $install:block, drop: |$drop_var:ident| $drop:block)? $(,)?}
+    ///      ),* $(,)?)
+    ///     => {
+    ///         $(
+    ///             profiler!(@inner { name: $name, port: $port $(,install: |$install_var| $install, drop: |$drop_var| $drop)? });
+    ///         )*
+    ///     };
+    /// 
+    ///     (@inner { name: $name:ident, port: $port:expr }) => {
+    ///         paste::paste!{
+    ///             #[doc = concat!("The address to bind the ", std::stringify!([< $name:lower >]), " thread profilers' server to")]
+    ///                 pub const [< $name:upper _PROFILER_ADDR >] : &'static str
+    ///                     = concat!("127.0.0.1:", $port);
+    ///
+    ///                 /// Installs the server's sink into the custom profiler
+    ///                 #[doc(hidden)]
+    ///                 fn [< $name:lower _profiler_server_install >](sink: puffin::FrameSink) -> puffin::FrameSinkId {
+    ///                     [< $name:lower _profiler_lock >]().add_sink(sink)
+    ///                 }
+    ///
+    ///                 /// Drops the server's sink and removes from profiler
+    ///                 #[doc(hidden)]
+    ///                 fn [< $name:lower _profiler_server_drop >](id: puffin::FrameSinkId){
+    ///                     [< $name:lower _profiler_lock >]().remove_sink(id);
+    ///                 }
+    ///
+    ///                 #[doc = concat!("The instance of the ", std::stringify!([< $name:lower >]), " thread profiler's server")]
+    ///                 pub static [< $name:upper _PROFILER_SERVER >] : once_cell::sync::Lazy<std::sync::Mutex<puffin_http::Server>>
+    ///                     = once_cell::sync::Lazy::new(|| {
+    ///                         eprintln!(
+    ///                             "starting puffin_http server for {} profiler at {}",
+    ///                             std::stringify!([<$name:lower>]),
+    ///                             [< $name:upper _PROFILER_ADDR >]
+    ///                         );
+    ///                         std::sync::Mutex::new(
+    ///                             puffin_http::Server::new_custom(
+    ///                                 [< $name:upper _PROFILER_ADDR >],
+    ///                                 // Can't use closures in a const context, use fn-pointers instead
+    ///                                 [< $name:lower _profiler_server_install >],
+    ///                                 [< $name:lower _profiler_server_drop >],
+    ///                             )
+    ///                             .expect(&format!("{} puffin_http server failed to start", std::stringify!([<$name:lower>])))
+    ///                         )
+    ///                     });
+    ///
+    ///                 #[doc = concat!("A custom reporter for the ", std::stringify!([< $name:lower >]), " thread reporter")]
+    ///                 pub fn [< $name:lower _profiler_reporter >] (info: puffin::ThreadInfo, details: &[puffin::ScopeDetails],  stream: &puffin::StreamInfoRef<'_>) {
+    ///                     [< $name:lower _profiler_lock >]().report(info, details, stream)
+    ///                 }
+    ///
+    ///                 #[doc = concat!("Accessor for the ", std::stringify!([< $name:lower >]), " thread reporter")]
+    ///                 pub fn [< $name:lower _profiler_lock >]() -> std::sync::MutexGuard<'static, puffin::GlobalProfiler> {
+    ///                     static [< $name _PROFILER >] : once_cell::sync::Lazy<std::sync::Mutex<puffin::GlobalProfiler>> = once_cell::sync::Lazy::new(Default::default);
+    ///                     [< $name _PROFILER >].lock().expect("poisoned std::sync::mutex")
+    ///                 }
+    ///
+    ///                 #[doc = concat!("Initialises the ", std::stringify!([< $name:lower >]), " thread reporter and server.\
+    ///                 Call this on each different thread you want to register with this profiler")]
+    ///                 pub fn [< $name:lower _profiler_init >]() {
+    ///                     eprintln!("init thread profiler \"{}\"", std::stringify!([<$name:lower>]));
+    ///                     std::mem::drop([< $name:upper _PROFILER_SERVER >].lock());
+    ///                     eprintln!("set thread custom profiler \"{}\"", std::stringify!([<$name:lower>]));
+    ///                     puffin::ThreadProfiler::initialize(::puffin::now_ns, [< $name:lower _profiler_reporter >]);
+    ///                 }
+    ///         }
+    ///     };
+    /// }
+    ///
+    /// profiler! {
+    ///     { name: UI,          port: "2a" },
+    ///     { name: RENDERER,    port: 8586 },
+    ///     { name: BACKGROUND,  port: 8587 },
+    /// }
+    /// 
+    /// pub fn demo() {
+    ///     std::thread::spawn(|| {
+    ///         // Initialise the custom profiler for this thread
+    ///         // Now all puffin events are sent to the custom profiling server instead
+    ///         //
+    ///         background_profiler_init();
+    ///
+    ///         for i in 0..100{
+    ///             puffin::profile_scope!("test");
+    ///             sleep(Duration::from_millis(i));
+    ///         }
+    ///
+    ///         // Mark a new frame so the data is flushed to the server
+    ///         background_profiler_lock().new_frame();
+    ///     });
+    /// }
+    /// ```
+    pub fn new_custom(
+        bind_addr: &str,
+        sink_install: fn (puffin::FrameSink) -> puffin::FrameSinkId,
+        sink_remove: fn (puffin::FrameSinkId) -> (),
+    ) -> anyhow::Result<Self> {
         let tcp_listener = TcpListener::bind(bind_addr).context("binding server TCP socket")?;
         tcp_listener
             .set_nonblocking(true)
@@ -65,14 +265,16 @@ impl Server {
             })
             .context("Couldn't spawn thread")?;
 
-        let sink_id = GlobalProfiler::lock().add_sink(Box::new(move |frame| {
+        // Call the `install` function to add ourselves as a sink
+        let sink_id = sink_install(Box::new(move |frame| {
             tx.send(frame).ok();
         }));
 
         Ok(Server {
             sink_id,
             join_handle: Some(join_handle),
             num_clients,
+            sink_remove,
         })
     }
 
@@ -84,7 +286,8 @@ impl Server {
 
 impl Drop for Server {
     fn drop(&mut self) {
-        GlobalProfiler::lock().remove_sink(self.sink_id);
+        // Remove ourselves from the profiler
+        (self.sink_remove)(self.sink_id);
 
         // Take care to send everything before we shut down:
         if let Some(join_handle) = self.join_handle.take() {