add `generic_activity_with_arg_recorder` to the self-profiler

This allows profiling costly arguments to be recorded only when `-Zself-profile-events=args` is on: using a closure that takes an `EventArgRecorder` and call its `record_arg` or `record_args` methods.

1 files changed, 88 insertions, 0 deletions

diff --git a/compiler/rustc_data_structures/src/profiling.rs b/compiler/rustc_data_structures/src/profiling.rs
index fa6c4a02cab..3f7a90a8467 100644
--- a/compiler/rustc_data_structures/src/profiling.rs
+++ b/compiler/rustc_data_structures/src/profiling.rs
@@ -97,6 +97,7 @@ use std::time::{Duration, Instant};
 pub use measureme::EventId;
 use measureme::{EventIdBuilder, Profiler, SerializableString, StringId};
 use parking_lot::RwLock;
+use smallvec::SmallVec;
 
 bitflags::bitflags! {
     struct EventFilter: u32 {
@@ -288,6 +289,66 @@ impl SelfProfilerRef {
         })
     }
 
+    /// Start profiling a generic activity, allowing costly arguments to be recorded. Profiling
+    /// continues until the `TimingGuard` returned from this call is dropped.
+    ///
+    /// If the arguments to a generic activity are cheap to create, use `generic_activity_with_arg`
+    /// or `generic_activity_with_args` for their simpler API. However, if they are costly or
+    /// require allocation in sufficiently hot contexts, then this allows for a closure to be called
+    /// only when arguments were asked to be recorded via `-Z self-profile-events=args`.
+    ///
+    /// In this case, the closure will be passed a `&mut EventArgRecorder`, to help with recording
+    /// one or many arguments within the generic activity being profiled, by calling its
+    /// `record_arg` method for example.
+    ///
+    /// This `EventArgRecorder` may implement more specific traits from other rustc crates, e.g. for
+    /// richer handling of rustc-specific argument types, while keeping this single entry-point API
+    /// for recording arguments.
+    ///
+    /// Note: recording at least one argument is *required* for the self-profiler to create the
+    /// `TimingGuard`. A panic will be triggered if that doesn't happen. This function exists
+    /// explicitly to record arguments, so it fails loudly when there are none to record.
+    ///
+    #[inline(always)]
+    pub fn generic_activity_with_arg_recorder<F>(
+        &self,
+        event_label: &'static str,
+        mut f: F,
+    ) -> TimingGuard<'_>
+    where
+        F: FnMut(&mut EventArgRecorder<'_>),
+    {
+        // Ensure this event will only be recorded when self-profiling is turned on.
+        self.exec(EventFilter::GENERIC_ACTIVITIES, |profiler| {
+            let builder = EventIdBuilder::new(&profiler.profiler);
+            let event_label = profiler.get_or_alloc_cached_string(event_label);
+
+            // Ensure the closure to create event arguments will only be called when argument
+            // recording is turned on.
+            let event_id = if profiler.event_filter_mask.contains(EventFilter::FUNCTION_ARGS) {
+                // Set up the builder and call the user-provided closure to record potentially
+                // costly event arguments.
+                let mut recorder = EventArgRecorder { profiler, args: SmallVec::new() };
+                f(&mut recorder);
+
+                // It is expected that the closure will record at least one argument. If that
+                // doesn't happen, it's a bug: we've been explicitly called in order to record
+                // arguments, so we fail loudly when there are none to record.
+                if recorder.args.is_empty() {
+                    panic!(
+                        "The closure passed to `generic_activity_with_arg_recorder` needs to \
+                         record at least one argument"
+                    );
+                }
+
+                builder.from_label_and_args(event_label, &recorder.args)
+            } else {
+                builder.from_label(event_label)
+            };
+            TimingGuard::start(profiler, profiler.generic_activity_event_kind, event_id)
+        })
+    }
+
     /// Record the size of an artifact that the compiler produces
     ///
     /// `artifact_kind` is the class of artifact (e.g., query_cache, object_file, etc.)
@@ -443,6 +504,33 @@ impl SelfProfilerRef {
     }
 }
 
+/// A helper for recording costly arguments to self-profiling events. Used with
+/// `SelfProfilerRef::generic_activity_with_arg_recorder`.
+pub struct EventArgRecorder<'p> {
+    /// The `SelfProfiler` used to intern the event arguments that users will ask to record.
+    profiler: &'p SelfProfiler,
+
+    /// The interned event arguments to be recorded in the generic activity event.
+    ///
+    /// The most common case, when actually recording event arguments, is to have one argument. Then
+    /// followed by recording two, in a couple places.
+    args: SmallVec<[StringId; 2]>,
+}
+
+impl EventArgRecorder<'_> {
+    /// Records a single argument within the current generic activity being profiled.
+    ///
+    /// Note: when self-profiling with costly event arguments, at least one argument
+    /// needs to be recorded. A panic will be triggered if that doesn't happen.
+    pub fn record_arg<A>(&mut self, event_arg: A)
+    where
+        A: Borrow<str> + Into<String>,
+    {
+        let event_arg = self.profiler.get_or_alloc_cached_string(event_arg);
+        self.args.push(event_arg);
+    }
+}
+
 pub struct SelfProfiler {
     profiler: Profiler,
     event_filter_mask: EventFilter,