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>

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();
+                            }
+                        });
+                    }
+                }
+            }
+        }
+    }
+}