Add tracing support

Signed-off-by: Will Killian <[email protected]>

Cargo.toml

@@ -25,16 +25,19 @@ publish.workspace = true
 license.workspace = true
 
 [features]
-default = ["color-name", "name-current-thread", "cuda", "cuda_runtime"]
+default = ["color-name", "name-current-thread", "cuda", "cuda_runtime", "tracing"]
 color-name = ["dep:color-name"]
 name-current-thread = ["dep:gettid"]
 cuda = ["nvtx-sys/cuda"]
 cuda_runtime = ["nvtx-sys/cuda_runtime"]
+tracing = ["dep:tracing", "dep:tracing-subscriber", "dep:color-name" ]
 
 [dependencies]
 color-name = { version = "1.1.0", optional = true }
 gettid = { version = "0.1.2", optional = true }
 nvtx-sys = { path = "crates/nvtx-sys" }
+tracing = { version = "0.1.40", optional = true, features = ["attributes"] }
+tracing-subscriber = { version = "0.3.18", optional = true }
 widestring = "1.0.2"
 
 [target.'cfg(unix)'.dependencies]
@@ -59,3 +62,7 @@ required-features = ["name-current-thread"]
 [[example]]
 name = "domain_threads"
 required-features = ["name-current-thread"]
+
+[[example]]
+name = "tracing"
+required-features = ["tracing"]

examples/tracing.rs

@@ -0,0 +1,35 @@
+use tracing::{info, instrument};
+use tracing_subscriber::prelude::*;
+
+#[instrument(fields(color = "goldenrod", category = 1, payload = i))]
+fn foo(i: u64) {
+    for j in 1..=i {
+        bar(j);
+        std::thread::sleep(std::time::Duration::from_millis(10 * i));
+    }
+}
+
+#[instrument(fields(color = "plum", category = 1, payload = j))]
+fn bar(j: u64) {
+    for k in 1..=j {
+        baz(k);
+        std::thread::sleep(std::time::Duration::from_millis(10 * j));
+    }
+}
+
+#[instrument(fields(color = "salmon", category = 2, payload = k))]
+fn baz(k: u64) {
+    std::thread::sleep(std::time::Duration::from_millis(10 * k));
+}
+
+fn main() {
+    let layer = nvtx::tracing::NvtxLayer::new("nvtx");
+    tracing_subscriber::registry().with(layer).init();
+
+    info!(
+        message = "At the beginning of the program",
+        color = "blue",
+        category = 2
+    );
+    foo(10);
+}

src/color.rs

@@ -99,7 +99,7 @@ impl TypeValueEncodable for Color {
     type Value = u32;
 
     fn encode(&self) -> (Self::Type, Self::Value) {
-        let as_u32 = u32::from_ne_bytes([self.a, self.r, self.g, self.b]);
+        let as_u32 = u32::from_be_bytes([self.a, self.r, self.g, self.b]);
         (Self::Type::NVTX_COLOR_ARGB, as_u32)
     }
 

src/domain/category.rs

@@ -3,6 +3,6 @@ use crate::Domain;
 /// Represents a category for use with event and range grouping. See [`Domain::register_category`], [`Domain::register_categories`]
 #[derive(Debug, Clone, Copy)]
 pub struct Category<'a> {
-    pub(super) id: u32,
-    pub(super) domain: &'a Domain,
+    pub(crate) id: u32,
+    pub(crate) domain: &'a Domain,
 }

src/domain/event_attributes.rs

@@ -4,10 +4,10 @@ use crate::{Color, Domain, Payload, TypeValueEncodable};
 /// All attributes that are associated with marks and ranges
 #[derive(Debug, Clone)]
 pub struct EventAttributes<'a> {
-    pub(super) category: Option<Category<'a>>,
-    pub(super) color: Option<Color>,
-    pub(super) payload: Option<Payload>,
-    pub(super) message: Option<Message<'a>>,
+    pub(crate) category: Option<Category<'a>>,
+    pub(crate) color: Option<Color>,
+    pub(crate) payload: Option<Payload>,
+    pub(crate) message: Option<Message<'a>>,
 }
 
 impl<'a> EventAttributes<'a> {

src/domain/local_range.rs

@@ -2,7 +2,7 @@ use super::EventArgument;
 use crate::Domain;
 use std::marker::PhantomData;
 
-/// A RAII-like object for modeling callstack Ranges within a Domain
+/// A RAII-like object for modeling callstack (thread-local) Ranges within a Domain
 #[derive(Debug)]
 pub struct LocalRange<'a> {
     pub(super) domain: &'a Domain,

src/domain/mod.rs

@@ -224,6 +224,20 @@ impl Domain {
         Range::new(arg, self)
     }
 
+    /// Internal function for starting a range and returning a raw Range Id
+    pub(crate) fn range_start<'a>(&self, arg: impl Into<EventArgument<'a>>) -> u64 {
+        let arg = match arg.into() {
+            EventArgument::Attributes(attr) => attr,
+            EventArgument::Message(m) => m.into(),
+        };
+        nvtx_sys::domain_range_start_ex(self.handle, &arg.encode())
+    }
+
+    /// Internal function for ending a range given a raw Range Id
+    pub(crate) fn range_end(&self, range_id: u64) {
+        nvtx_sys::domain_range_end(self.handle, range_id);
+    }
+
     /// Name a resource
     ///
     /// ```

src/domain/range.rs

@@ -1,7 +1,7 @@
 use super::EventArgument;
 use crate::Domain;
 
-/// A RAII-like object for modeling start/end Ranges within a Domain
+/// A RAII-like object for modeling process-wide Ranges within a Domain
 #[derive(Debug)]
 pub struct Range<'a> {
     pub(super) id: nvtx_sys::RangeId,
@@ -10,18 +10,16 @@ pub struct Range<'a> {
 
 impl<'a> Range<'a> {
     pub(super) fn new(arg: impl Into<EventArgument<'a>>, domain: &'a Domain) -> Range<'a> {
-        let arg = match arg.into() {
-            EventArgument::Attributes(attr) => attr,
-            EventArgument::Message(m) => m.into(),
-        };
-        let id = nvtx_sys::domain_range_start_ex(domain.handle, &arg.encode());
-        Range { id, domain }
+        Range {
+            id: domain.range_start(arg),
+            domain,
+        }
     }
 }
 
 impl<'a> Drop for Range<'a> {
     fn drop(&mut self) {
-        nvtx_sys::domain_range_end(self.domain.handle, self.id)
+        self.domain.range_end(self.id);
     }
 }
 

src/lib.rs

@@ -38,6 +38,9 @@
 //!   runtime resources such as Devices, Events, and Streams. The feature also adds `CudaRuntimeIdentifier`
 //!   to the [`crate::domain`] module to provide an alternative naming mechanism via [`crate::Domain::name_resource`].
 //!
+//! * **tracing** -
+//!   When enabled,
+//!
 //! ## Platform-specific types
 //!
 //! * **PThread Resource Naming** -
@@ -46,12 +49,17 @@
 
 /// color support
 pub mod color;
+
 /// specialized types for use within a domain context
 pub mod domain;
 
 /// platform native types
 pub mod native_types;
 
+#[cfg(feature = "tracing")]
+/// tracing support
+pub mod tracing;
+
 mod category;
 mod event_argument;
 mod event_attributes;

src/local_range.rs

@@ -1,7 +1,7 @@
 use crate::{EventArgument, Message};
 use std::marker::PhantomData;
 
-/// A RAII-like object for modeling callstack Ranges
+/// A RAII-like object for modeling callstack (thread-local) Ranges
 #[derive(Debug)]
 pub struct LocalRange {
     // prevent Sync + Send

src/range.rs

@@ -1,6 +1,6 @@
 use crate::{EventArgument, Message};
 
-/// A RAII-like object for modeling start/end Ranges
+/// A RAII-like object for modeling process-wide Ranges
 #[derive(Debug)]
 pub struct Range {
     id: nvtx_sys::RangeId,