1 files changed, 100 insertions, 0 deletions
diff --git a/profiling/src/lib.rs b/profiling/src/lib.rs
new file mode 100644
index 00000000..8fc58fb8
--- /dev/null
+++ b/profiling/src/lib.rs
@@ -0,0 +1,100 @@
+use std::time::Duration;
+use tracing_subscriber::prelude::*;
+use tracing_subscriber::Registry;
+
+#[cfg(feature = "tracing-chrome")]
+use {
+    tracing_chrome::FlushGuard,
+    tracing_subscriber::fmt::{format::DefaultFields, FormattedFields},
+};
+
+pub use tracing::{info_span, instrument};
+
+/// Profiler state. This will likely need to be updated or reworked when adding new tracing backends.
+pub struct Profiler {
+    #[cfg(feature = "tracing-chrome")]
+    /// [`FlushGuard`] must not be dropped until the application scope is dropped for accurate tracing.
+    _guard: FlushGuard,
+}
+
+pub fn init() -> Profiler {
+    // Registry stores the spans & generates unique span IDs
+    let subscriber = Registry::default();
+
+    #[cfg(feature = "tracing-chrome")]
+    let (chrome_layer, guard) = {
+        let mut layer = tracing_chrome::ChromeLayerBuilder::new();
+
+        // Optional configurable env var: CHROME_TRACE_FILE=/path/to/trace_file/file.json,
+        // for uploading to chrome://tracing (old) or ui.perfetto.dev (new).
+        if let Ok(path) = std::env::var("CHROME_TRACE_FILE") {
+            layer = layer.file(path);
+        } else if let Ok(current_dir) = std::env::current_dir() {
+            let time = std::time::SystemTime::now()
+                .duration_since(std::time::UNIX_EPOCH)
+                .unwrap_or(Duration::from_millis(0))
+                .as_millis();
+
+            let trace_file_name = current_dir
+                .file_name()
+                .map(|file_dir| {
+                    format!(
+                        "{}_trace_{}.json",
+                        file_dir.to_str().unwrap_or("trace"),
+                        time
+                    )
+                })
+                .unwrap_or_else(|| "trace.json".to_string());
+
+            let path = format!(
+                "{}/{}",
+                current_dir.to_str().expect("Invalid path"),
+                trace_file_name
+            );
+
+            layer = layer.file(path);
+        } else {
+            layer = layer.file(env!("CARGO_MANIFEST_DIR"))
+        }
+
+        let (chrome_layer, guard) = layer
+            .name_fn(Box::new(|event_or_span| match event_or_span {
+                tracing_chrome::EventOrSpan::Event(event) => {
+                    event.metadata().name().into()
+                }
+                tracing_chrome::EventOrSpan::Span(span) => {
+                    if let Some(fields) = span
+                        .extensions()
+                        .get::<FormattedFields<DefaultFields>>()
+                    {
+                        format!(
+                            "{}: {}",
+                            span.metadata().name(),
+                            fields.fields.as_str()
+                        )
+                    } else {
+                        span.metadata().name().into()
+                    }
+                }
+            }))
+            .build();
+
+        (chrome_layer, guard)
+    };
+
+    let fmt_layer = tracing_subscriber::fmt::Layer::default();
+    let subscriber = subscriber.with(fmt_layer);
+
+    #[cfg(feature = "tracing-chrome")]
+    let subscriber = subscriber.with(chrome_layer);
+
+    // create dispatcher which will forward span events to the subscriber
+    // this can only be set once or will panic
+    tracing::subscriber::set_global_default(subscriber)
+        .expect("Profiler could not set the global default subscriber.");
+
+    Profiler {
+        #[cfg(feature = "tracing-chrome")]
+        _guard: guard,
+    }
+}