Plumb selected tracing spans into hitrace (#33324)

* Plumb selected tracing spans into hitrace

Signed-off-by: Delan Azabani <dazabani@igalia.com>

* Tag the spans in #33189 with the `servo_profiling` field

Signed-off-by: Delan Azabani <dazabani@igalia.com>

* Add comment about fields vs extensions

Signed-off-by: Delan Azabani <dazabani@igalia.com>

---------

Signed-off-by: Delan Azabani <dazabani@igalia.com>

7 files changed, 162 insertions, 11 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 754f6fdf97f..9980ee4df4c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2875,6 +2875,21 @@ dependencies = [
 ]
 
 [[package]]
+name = "hitrace"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f92c0ae6f30b32eaeb811143fba3b56864f477b2e69458b13779a07b3aaf2f6e"
+dependencies = [
+ "hitrace-sys",
+]
+
+[[package]]
+name = "hitrace-sys"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "315f4e893d1caac3a97b9e6cbcf211a7012c6615cd688e4430f0cd5712ac3a3f"
+
+[[package]]
 name = "home"
 version = "0.5.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -6289,6 +6304,7 @@ dependencies = [
  "glow 0.13.1",
  "headers",
  "hilog",
+ "hitrace",
  "http",
  "image",
  "ipc-channel",
diff --git a/Cargo.toml b/Cargo.toml
index 8d0b85bfb1a..076cdbd74d0 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -58,6 +58,7 @@ gstreamer-gl-sys = "0.22"
 gstreamer-sys = "0.22"
 gstreamer-video = "0.22"
 headers = "0.3"
+hitrace = "0.1.4"
 html5ever = "0.28"
 http = "0.2"
 hyper = "0.14"
diff --git a/components/compositing/compositor.rs b/components/compositing/compositor.rs
index 2b072481f9a..02ed570ae2c 100644
--- a/components/compositing/compositor.rs
+++ b/components/compositing/compositor.rs
@@ -682,7 +682,7 @@ impl<Window: WindowMethods + ?Sized> IOCompositor<Window> {
 
     /// Accept messages from content processes that need to be relayed to the WebRender
     /// instance in the parent process.
-    #[tracing::instrument(skip(self))]
+    #[tracing::instrument(skip(self), fields(servo_profiling = true))]
     fn handle_webrender_message(&mut self, msg: ForwardedToCompositorMsg) {
         match msg {
             ForwardedToCompositorMsg::Layout(ScriptToCompositorMsg::SendInitialTransaction(
@@ -770,7 +770,11 @@ impl<Window: WindowMethods + ?Sized> IOCompositor<Window> {
                     display_list_descriptor,
                 );
 
-                let span = span!(Level::TRACE, "ScriptToCompositorMsg::BuiltDisplayList");
+                let span = span!(
+                    Level::TRACE,
+                    "ScriptToCompositorMsg::BuiltDisplayList",
+                    servo_profiling = true
+                );
                 let _enter = span.enter();
                 let pipeline_id = display_list_info.pipeline_id;
                 let details = self.pipeline_details(pipeline_id.into());
@@ -2058,7 +2062,7 @@ impl<Window: WindowMethods + ?Sized> IOCompositor<Window> {
     /// Returns Ok if composition was performed or Err if it was not possible to composite for some
     /// reason. When the target is [CompositeTarget::SharedMemory], the image is read back from the
     /// GPU and returned as Ok(Some(png::Image)), otherwise we return Ok(None).
-    #[tracing::instrument(skip(self))]
+    #[tracing::instrument(skip(self), fields(servo_profiling = true))]
     fn composite_specific_target(
         &mut self,
         target: CompositeTarget,
@@ -2271,7 +2275,11 @@ impl<Window: WindowMethods + ?Sized> IOCompositor<Window> {
             },
         };
 
-        let span = span!(Level::TRACE, "ConstellationMsg::ReadyToPresent");
+        let span = span!(
+            Level::TRACE,
+            "ConstellationMsg::ReadyToPresent",
+            servo_profiling = true
+        );
         let _enter = span.enter();
         // Notify embedder that servo is ready to present.
         // Embedder should call `present` to tell compositor to continue rendering.
@@ -2297,9 +2305,13 @@ impl<Window: WindowMethods + ?Sized> IOCompositor<Window> {
             .map(|info| info.framebuffer_id())
     }
 
-    #[tracing::instrument(skip(self))]
+    #[tracing::instrument(skip(self), fields(servo_profiling = true))]
     pub fn present(&mut self) {
-        let span = span!(Level::TRACE, "Compositor Present Surface");
+        let span = span!(
+            Level::TRACE,
+            "Compositor Present Surface",
+            servo_profiling = true
+        );
         let _enter = span.enter();
         if let Err(err) = self.rendering_context.present() {
             warn!("Failed to present surface: {:?}", err);
@@ -2362,7 +2374,7 @@ impl<Window: WindowMethods + ?Sized> IOCompositor<Window> {
         );
     }
 
-    #[tracing::instrument(skip(self))]
+    #[tracing::instrument(skip(self), fields(servo_profiling = true))]
     pub fn receive_messages(&mut self) -> bool {
         // Check for new messages coming from the other threads in the system.
         let mut compositor_messages = vec![];
@@ -2389,7 +2401,7 @@ impl<Window: WindowMethods + ?Sized> IOCompositor<Window> {
         true
     }
 
-    #[tracing::instrument(skip(self))]
+    #[tracing::instrument(skip(self), fields(servo_profiling = true))]
     pub fn perform_updates(&mut self) -> bool {
         if self.shutdown_state == ShutdownState::FinishedShuttingDown {
             return false;
diff --git a/components/constellation/constellation.rs b/components/constellation/constellation.rs
index f353616cf75..24bcfd0906c 100644
--- a/components/constellation/constellation.rs
+++ b/components/constellation/constellation.rs
@@ -1333,7 +1333,7 @@ where
         }
     }
 
-    #[tracing::instrument(skip(self))]
+    #[tracing::instrument(skip(self), fields(servo_profiling = true))]
     fn handle_request_from_compositor(&mut self, message: FromCompositorMsg) {
         trace_msg_from_compositor!(message, "{message:?}");
         match message {
@@ -1546,7 +1546,11 @@ where
                 self.set_webview_throttled(webview_id, throttled);
             },
             FromCompositorMsg::ReadyToPresent(webview_ids) => {
-                let span = span!(Level::TRACE, "FromCompositorMsg::ReadyToPresent");
+                let span = span!(
+                    Level::TRACE,
+                    "FromCompositorMsg::ReadyToPresent",
+                    servo_profiling = true
+                );
                 let _enter = span.enter();
                 self.embedder_proxy
                     .send((None, EmbedderMsg::ReadyToPresent(webview_ids)));
diff --git a/components/servo/lib.rs b/components/servo/lib.rs
index ed91d0947dc..10b701eae93 100644
--- a/components/servo/lib.rs
+++ b/components/servo/lib.rs
@@ -225,7 +225,7 @@ impl<Window> Servo<Window>
 where
     Window: WindowMethods + 'static + ?Sized,
 {
-    #[tracing::instrument(skip(embedder, window))]
+    #[tracing::instrument(skip(embedder, window), fields(servo_profiling = true))]
     pub fn new(
         mut embedder: Box<dyn EmbedderMethods>,
         window: Rc<Window>,
diff --git a/ports/servoshell/Cargo.toml b/ports/servoshell/Cargo.toml
index cc5047f344c..172c4f30262 100644
--- a/ports/servoshell/Cargo.toml
+++ b/ports/servoshell/Cargo.toml
@@ -49,6 +49,7 @@ native-bluetooth = ["libservo/native-bluetooth"]
 profilemozjs = ["libservo/profilemozjs"]
 refcell_backtrace = ["libservo/refcell_backtrace"]
 tracing = ["dep:tracing", "dep:tracing-subscriber"]
+tracing-hitrace = ["tracing", "dep:hitrace"]
 tracing-perfetto = ["tracing", "dep:tracing-perfetto"]
 webdriver = ["libservo/webdriver"]
 webgl_backtrace = ["libservo/webgl_backtrace"]
@@ -59,6 +60,7 @@ libservo = { path = "../../components/servo" }
 cfg-if = { workspace = true }
 log = { workspace = true }
 getopts = { workspace = true }
+hitrace = { workspace = true, optional = true }
 mime_guess = { workspace = true }
 url = { workspace = true }
 servo-media = { workspace = true }
diff --git a/ports/servoshell/lib.rs b/ports/servoshell/lib.rs
index 6120749c6cf..5fd4b9fbba0 100644
--- a/ports/servoshell/lib.rs
+++ b/ports/servoshell/lib.rs
@@ -2,6 +2,8 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
 
+use cfg_if::cfg_if;
+
 #[cfg(any(target_os = "macos", target_os = "linux"))]
 #[macro_use]
 extern crate sig;
@@ -54,6 +56,13 @@ pub fn init_tracing() {
             subscriber.with(perfetto_layer)
         };
 
+        #[cfg(feature = "tracing-hitrace")]
+        let subscriber = {
+            use tracing_subscriber::layer::SubscriberExt;
+            // Set up a HitraceLayer for performance tracing.
+            subscriber.with(HitraceLayer::default())
+        };
+
         // Same as SubscriberInitExt::init, but avoids initialising the tracing-log compat layer,
         // since it would break Servo’s FromScriptLogger and FromCompositorLogger.
         // <https://docs.rs/tracing-subscriber/0.3.18/tracing_subscriber/util/trait.SubscriberInitExt.html#method.init>
@@ -70,3 +79,110 @@ pub fn servo_version() -> String {
         env!("VERGEN_GIT_SHA")
     )
 }
+
+/// Plumbs tracing spans into HiTrace, with the following caveats:
+///
+/// - We ignore spans unless they have a `servo_profiling` field.
+/// - We map span entry ([`Layer::on_enter`]) to `OH_HiTrace_StartTrace(metadata.name())`.
+/// - We map span exit ([`Layer::on_exit`]) to `OH_HiTrace_FinishTrace()`.
+///
+/// As a result, within each thread, spans must exit in reverse order of their entry, otherwise the
+/// resultant profiling data will be incorrect (see the section below). This is not necessarily the
+/// case for tracing spans, since there can be multiple [trace trees], so we check that this
+/// invariant is upheld when debug assertions are enabled, logging errors if it is violated.
+///
+/// [trace trees]: https://docs.rs/tracing/0.1.40/tracing/span/index.html#span-relationships
+///
+/// # Uniquely identifying spans
+///
+/// We need to ensure that the start and end points of one span are not mixed up with other spans.
+/// For now, we use the HiTrace [synchronous API], which restricts how spans must behave.
+///
+/// In the HiTrace [synchronous API], spans must have stack-like behaviour, because spans are keyed
+/// entirely on their *name* string, and OH_HiTrace_FinishTrace always ends the most recent span.
+/// While synchronous API spans are thread-local, callers could still violate this invariant with
+/// reentrant or asynchronous code.
+///
+/// In the [asynchronous API], spans are keyed on a (*name*,*taskId*) pair, where *name* is again
+/// a string, and *taskId* is an arbitrary [`i32`]. This makes *taskId* a good place for a unique
+/// identifier, but asynchronous spans can cross thread boundaries, so the identifier needs to be
+/// temporally unique in the whole process.
+///
+/// Tracing spans have such an identifier ([`Id`]), but they’re [`u64`]-based, and their format
+/// is an internal implementation detail of the [`Subscriber`]. For [`Registry`], those values
+/// [come from] a [packed representation] of a generation number, thread number, page number, and
+/// variable-length index. This makes them hard to compress robustly into an [`i32`].
+///
+/// If we move to the asynchronous API, we will need to generate our own *taskId* values, perhaps
+/// by combining some sort of thread id with a thread-local atomic counter. [`ThreadId`] is opaque
+/// in stable Rust, and converts to a [`u64`] in unstable Rust, so we would also need to make our
+/// own thread ids, perhaps by having a global atomic counter cached in a thread-local.
+///
+/// [synchronous API]: https://docs.rs/hitrace-sys/0.1.2/hitrace_sys/fn.OH_HiTrace_StartTrace.html
+/// [asynchronous API]: https://docs.rs/hitrace-sys/0.1.2/hitrace_sys/fn.OH_HiTrace_StartAsyncTrace.html
+/// [`Registry`]: tracing_subscriber::Registry
+/// [come from]: https://docs.rs/tracing-subscriber/0.3.18/src/tracing_subscriber/registry/sharded.rs.html#237-269
+/// [packed representation]: https://docs.rs/sharded-slab/0.1.7/sharded_slab/trait.Config.html
+/// [`ThreadId`]: std::thread::ThreadId
+#[cfg(feature = "tracing-hitrace")]
+#[derive(Default)]
+struct HitraceLayer {}
+
+cfg_if! {
+    if #[cfg(feature = "tracing-hitrace")] {
+        use std::cell::RefCell;
+
+        use tracing::span::Id;
+        use tracing::Subscriber;
+        use tracing_subscriber::Layer;
+
+        #[cfg(debug_assertions)]
+        thread_local! {
+            /// Stack of span names, to ensure the HiTrace synchronous API is not misused.
+            static HITRACE_NAME_STACK: RefCell<Vec<String>> = RefCell::default();
+        }
+
+        impl<S: Subscriber + for<'lookup> tracing_subscriber::registry::LookupSpan<'lookup>>
+            Layer<S> for HitraceLayer
+        {
+            fn on_enter(&self, id: &Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
+                if let Some(metadata) = ctx.metadata(id) {
+                    // TODO: is this expensive? Would extensions be faster?
+                    // <https://docs.rs/tracing-subscriber/0.3.18/tracing_subscriber/registry/struct.ExtensionsMut.html>
+                    if metadata.fields().field("servo_profiling").is_some() {
+                        #[cfg(debug_assertions)]
+                        HITRACE_NAME_STACK.with_borrow_mut(|stack|
+                            stack.push(metadata.name().to_owned()));
+
+                        hitrace::start_trace(
+                            &std::ffi::CString::new(metadata.name())
+                                .expect("Failed to convert str to CString"),
+                        );
+                    }
+                }
+            }
+
+            fn on_exit(&self, id: &Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
+                if let Some(metadata) = ctx.metadata(id) {
+                    if metadata.fields().field("servo_profiling").is_some() {
+                        hitrace::finish_trace();
+
+                        #[cfg(debug_assertions)]
+                        HITRACE_NAME_STACK.with_borrow_mut(|stack| {
+                            if stack.last().map(|name| &**name) != Some(metadata.name()) {
+                                log::error!(
+                                    "Tracing span out of order: {} (stack: {:?})",
+                                    metadata.name(),
+                                    stack
+                                );
+                            }
+                            if !stack.is_empty() {
+                                stack.pop();
+                            }
+                        });
+                    }
+                }
+            }
+        }
+    }
+}