Speed up the subscriber

Better string handling, avoiding work where possible, new filter mode to
optimize the case when the subscriber is the only layer.

Bug: 293435349
Test: atest
Test: benchmark
Test: manually ran the sample and captured a trace
Change-Id: I05ede25daaadb109cffae2ad7658d4fab58b4b5d

3 files changed, 243 insertions, 70 deletions

diff --git a/libatrace_rust/benchmark/README.md b/libatrace_rust/benchmark/README.md
index e8532952..7c9e64c1 100644
--- a/libatrace_rust/benchmark/README.md
+++ b/libatrace_rust/benchmark/README.md
@@ -22,7 +22,8 @@ There is a benchmark for the tracing crate subscriber - `libatrace_tracing_subsc
 We use it to check overhead over the base `libatrace_rust`.
 
 Similarly to the wrapper benchmarks, the subscriber is measured with tracing off and on. There are
-cases with and without extra fields to measure the cost of formatting.
+cases with and without extra fields to measure the cost of formatting. Cases that start with
+`filtered_` measure the subscriber in filtering mode with tracing disabled.
 
 ## Running the benchmarks
 
@@ -83,12 +84,16 @@ the subscriber itself.
 Results from `libatrace_tracing_subscriber_benchmark 2>&1 | grep time`:
 
 ```text
-tracing_off_event       time:   [298.67 ns 299.99 ns 301.52 ns]
-tracing_off_event_args  time:   [2.1225 µs 2.1429 µs 2.1647 µs]
-tracing_off_span        time:   [610.88 ns 619.88 ns 628.70 ns]
-tracing_off_span_args   time:   [2.6585 µs 2.6698 µs 2.6845 µs]
-tracing_on_event        time:   [1.6249 µs 1.6495 µs 1.6887 µs]
-tracing_on_event_args   time:   [3.5608 µs 3.5831 µs 3.6064 µs]
-tracing_on_span         time:   [3.1778 µs 3.1982 µs 3.2211 µs]
-tracing_on_span_args    time:   [5.4340 µs 5.4717 µs 5.5143 µs]
+tracing_off_event       time:   [47.444 ns 47.945 ns 48.585 ns]
+filtered_event          time:   [26.852 ns 26.942 ns 27.040 ns]
+tracing_off_event_args  time:   [80.597 ns 80.997 ns 81.475 ns]
+filtered_event_args     time:   [26.680 ns 26.782 ns 26.887 ns]
+tracing_off_span        time:   [316.48 ns 317.72 ns 319.12 ns]
+filtered_span           time:   [27.900 ns 27.959 ns 28.018 ns]
+tracing_off_span_args   time:   [364.92 ns 367.57 ns 370.95 ns]
+filtered_span_args      time:   [27.625 ns 27.919 ns 28.207 ns]
+tracing_on_event        time:   [1.4639 µs 1.4805 µs 1.4954 µs]
+tracing_on_event_args   time:   [2.0088 µs 2.0197 µs 2.0314 µs]
+tracing_on_span         time:   [2.7907 µs 2.7996 µs 2.8103 µs]
+tracing_on_span_args    time:   [3.6846 µs 3.6992 µs 3.7168 µs]
 ```
diff --git a/libatrace_rust/benchmark/src/tracing_subscriber_benchmark.rs b/libatrace_rust/benchmark/src/tracing_subscriber_benchmark.rs
index 17e9e44a..f5a3a3b4 100644
--- a/libatrace_rust/benchmark/src/tracing_subscriber_benchmark.rs
+++ b/libatrace_rust/benchmark/src/tracing_subscriber_benchmark.rs
@@ -33,11 +33,26 @@ where
     });
 }
 
+fn bench_with_filtering_subscriber<F>(c: &mut Criterion, name: &str, mut f: F)
+where
+    F: FnMut(),
+{
+    let subscriber = tracing_subscriber::registry().with(AtraceSubscriber::default().with_filter());
+    tracing::subscriber::with_default(subscriber, || {
+        c.bench_function(name, |b| b.iter(&mut f));
+    });
+}
+
 fn bench_tracing_off_event(c: &mut Criterion) {
     turn_tracing_off();
     bench_with_subscriber(c, "tracing_off_event", || tracing::info!("bench info event"));
 }
 
+fn bench_filtered_event(c: &mut Criterion) {
+    turn_tracing_off();
+    bench_with_filtering_subscriber(c, "filtered_event", || tracing::info!("bench info event"));
+}
+
 fn bench_tracing_off_event_args(c: &mut Criterion) {
     turn_tracing_off();
     let v = make_example_vec();
@@ -50,6 +65,18 @@ fn bench_tracing_off_event_args(c: &mut Criterion) {
     });
 }
 
+fn bench_filtered_event_args(c: &mut Criterion) {
+    turn_tracing_off();
+    let v = make_example_vec();
+    bench_with_filtering_subscriber(c, "filtered_event_args", || {
+        tracing::info!(debug_arg1 = 123,
+            debug_arg2 = "argument",
+            debug_arg3 = ?v,
+            debug_arg4 = "last",
+            "bench info event")
+    });
+}
+
 fn bench_tracing_off_span(c: &mut Criterion) {
     turn_tracing_off();
     bench_with_subscriber(c, "tracing_off_span", || {
@@ -57,6 +84,13 @@ fn bench_tracing_off_span(c: &mut Criterion) {
     });
 }
 
+fn bench_filtered_span(c: &mut Criterion) {
+    turn_tracing_off();
+    bench_with_filtering_subscriber(c, "filtered_span", || {
+        let _entered = tracing::info_span!("bench info span").entered();
+    });
+}
+
 fn bench_tracing_off_span_args(c: &mut Criterion) {
     turn_tracing_off();
     let v = make_example_vec();
@@ -69,6 +103,18 @@ fn bench_tracing_off_span_args(c: &mut Criterion) {
     });
 }
 
+fn bench_filtered_span_args(c: &mut Criterion) {
+    turn_tracing_off();
+    let v = make_example_vec();
+    bench_with_filtering_subscriber(c, "filtered_span_args", || {
+        let _entered = tracing::info_span!("bench info span", debug_arg1 = 123,
+            debug_arg2 = "argument",
+            debug_arg3 = ?v,
+            debug_arg4 = "last")
+        .entered();
+    });
+}
+
 fn bench_tracing_on_event(c: &mut Criterion) {
     turn_tracing_on();
     bench_with_subscriber(c, "tracing_on_event", || tracing::info!("bench info event"));
@@ -113,9 +159,13 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     let mut criterion = new_criterion();
 
     bench_tracing_off_event(&mut criterion);
+    bench_filtered_event(&mut criterion);
     bench_tracing_off_event_args(&mut criterion);
+    bench_filtered_event_args(&mut criterion);
     bench_tracing_off_span(&mut criterion);
+    bench_filtered_span(&mut criterion);
     bench_tracing_off_span_args(&mut criterion);
+    bench_filtered_span_args(&mut criterion);
 
     bench_tracing_on_event(&mut criterion);
     bench_tracing_on_event_args(&mut criterion);
diff --git a/libatrace_rust/src/tracing_subscriber.rs b/libatrace_rust/src/tracing_subscriber.rs
index 06bec9a8..76c20740 100644
--- a/libatrace_rust/src/tracing_subscriber.rs
+++ b/libatrace_rust/src/tracing_subscriber.rs
@@ -15,18 +15,19 @@
 //! Tracing-subscriber layer for libatrace_rust.
 
 use ::atrace::AtraceTag;
-use tracing::{Event, Id, Subscriber};
-use tracing_subscriber::layer::{Context, Layer};
-use tracing_subscriber::registry::LookupSpan;
-
+use std::fmt::Write;
 use tracing::span::Attributes;
 use tracing::span::Record;
+use tracing::{Event, Id, Subscriber};
 use tracing_subscriber::field::Visit;
+use tracing_subscriber::layer::{Context, Layer};
+use tracing_subscriber::registry::LookupSpan;
 
 /// Subscriber layer that forwards events to ATrace.
 pub struct AtraceSubscriber {
     tag: AtraceTag,
     should_record_fields: bool,
+    should_filter: bool,
 }
 
 impl Default for AtraceSubscriber {
@@ -38,113 +39,193 @@ impl Default for AtraceSubscriber {
 impl AtraceSubscriber {
     /// Makes a new subscriber with tag.
     pub fn new(tag: AtraceTag) -> AtraceSubscriber {
-        AtraceSubscriber { tag, should_record_fields: true }
+        AtraceSubscriber { tag, should_filter: false, should_record_fields: true }
     }
 
-    /// Disables recording of field values.
-    pub fn without_fields(self) -> AtraceSubscriber {
-        AtraceSubscriber { tag: self.tag, should_record_fields: false }
+    /// Enables event and span filtering. With filtering enabled, this layer would filter events for
+    /// all the layers of the subscriber.
+    /// Use this to speed up the subscriber if it's the only layer. Do not enable if you need other
+    /// layers to receive events when ATrace is disabled.
+    pub fn with_filter(self) -> AtraceSubscriber {
+        AtraceSubscriber { should_filter: true, ..self }
     }
-}
 
-struct FieldFormatter {
-    is_event: bool,
-    s: String,
-}
-
-impl FieldFormatter {
-    fn for_event() -> FieldFormatter {
-        FieldFormatter { is_event: true, s: String::new() }
-    }
-    fn for_span() -> FieldFormatter {
-        FieldFormatter { is_event: false, s: String::new() }
+    /// Disables recording of field values.
+    pub fn without_fields(self) -> AtraceSubscriber {
+        AtraceSubscriber { should_record_fields: false, ..self }
     }
 }
 
-impl FieldFormatter {
-    fn as_string(&self) -> &String {
-        &self.s
-    }
-    fn is_empty(&self) -> bool {
-        self.s.is_empty()
+// Internal methods.
+impl AtraceSubscriber {
+    /// Checks that events and spans should be recorded in the span/event notification.
+    fn should_process_event(&self) -> bool {
+        // If `should_filter == true` we don't need to check the tag - it was already checked by
+        // the layer filter in the `Layer::enabled()` method.
+        // The checks are done in this order:
+        //  * `Layer::register_callsite()` - once per callsite, the result is cached.
+        //  * `Layer::enabled()` - once per span or event construction if the callsite is enabled.
+        //  * `should_process_event()` - on every notification like new span, span enter/exit/record, event.
+        // The first two checks are global, i.e. affect other layers, and only enabled with `should_filter`.
+        // Read more:
+        // https://docs.rs/tracing-subscriber/latest/tracing_subscriber/layer/index.html#filtering-with-layers
+        self.should_filter || atrace::atrace_is_tag_enabled(self.tag)
     }
 }
 
-impl Visit for FieldFormatter {
-    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
-        if !self.s.is_empty() {
-            self.s.push_str(", ");
-        }
-        if self.is_event && field.name() == "message" {
-            self.s.push_str(&format!("{:?}", value));
+impl<S: Subscriber + for<'lookup> LookupSpan<'lookup>> Layer<S> for AtraceSubscriber {
+    fn register_callsite(
+        &self,
+        _metadata: &'static tracing::Metadata<'static>,
+    ) -> tracing::subscriber::Interest {
+        if self.should_filter {
+            // When we return `Interest::sometimes()`, the `enabled()` method would get checked
+            // every time.
+            // We can't use callsite caching (`Interest::never()`) because there's no callback
+            // for when tracing gets enabled - we need to check it every time.
+            tracing::subscriber::Interest::sometimes()
         } else {
-            self.s.push_str(&format!("{} = {:?}", field.name(), value));
+            // If we do not disable events in the layer, we always receive the notifications.
+            tracing::subscriber::Interest::always()
         }
     }
-}
 
-impl<S: Subscriber + for<'lookup> LookupSpan<'lookup>> Layer<S> for AtraceSubscriber {
+    // When filtering in this layer is enabled, this method would get called on every event and span.
+    // This filter affects all layers, so if this method returns false, it would disable the event
+    // for others as well.
+    fn enabled(&self, _metadata: &tracing::Metadata<'_>, _ctx: Context<'_, S>) -> bool {
+        !self.should_filter || atrace::atrace_is_tag_enabled(self.tag)
+    }
+
     fn on_new_span(&self, attrs: &Attributes, id: &Id, ctx: Context<S>) {
-        if self.should_record_fields {
-            let mut formatter = FieldFormatter::for_span();
-            attrs.record(&mut formatter);
-            ctx.span(id).unwrap().extensions_mut().insert(formatter)
+        if !self.should_record_fields || attrs.fields().is_empty() || !self.should_process_event() {
+            return;
         }
+
+        let span = ctx.span(id).unwrap();
+        let mut formatter = FieldFormatter::for_span(span.metadata().name());
+        attrs.record(&mut formatter);
+        span.extensions_mut().insert(formatter);
     }
 
     fn on_record(&self, span: &Id, values: &Record, ctx: Context<S>) {
-        if self.should_record_fields {
-            values.record(
-                ctx.span(span).unwrap().extensions_mut().get_mut::<FieldFormatter>().unwrap(),
-            );
+        if !self.should_record_fields || !self.should_process_event() {
+            return;
         }
+
+        values
+            .record(ctx.span(span).unwrap().extensions_mut().get_mut::<FieldFormatter>().unwrap());
     }
 
     fn on_enter(&self, id: &Id, ctx: Context<S>) {
+        if !self.should_process_event() {
+            return;
+        }
+
         let span = ctx.span(id).unwrap();
-        let mut span_str = String::from(span.metadata().name());
-        if self.should_record_fields {
+        if span.fields().is_empty() || !self.should_record_fields {
+            atrace::atrace_begin(self.tag, span.metadata().name());
+        } else {
             let span_extensions = span.extensions();
             let formatter = span_extensions.get::<FieldFormatter>().unwrap();
-            if !formatter.is_empty() {
-                span_str.push_str(", ");
-                span_str.push_str(formatter.as_string());
-            }
+            atrace::atrace_begin(self.tag, formatter.as_str());
         }
-        atrace::atrace_begin(self.tag, &span_str);
     }
 
     fn on_exit(&self, _id: &Id, _ctx: Context<S>) {
+        if !self.should_process_event() {
+            return;
+        }
+
         atrace::atrace_end(self.tag);
     }
 
     fn on_event(&self, event: &Event, _ctx: Context<S>) {
-        let mut event_str = String::new();
+        if !self.should_process_event() {
+            return;
+        }
+
         if self.should_record_fields {
             let mut formatter = FieldFormatter::for_event();
             event.record(&mut formatter);
-            event_str = formatter.as_string().clone();
-        } else {
+            atrace::atrace_instant(self.tag, formatter.as_str());
+        } else if let Some(field) = event.metadata().fields().field("message") {
             struct MessageVisitor<'a> {
-                s: &'a mut String,
+                tag: AtraceTag,
+                field: &'a tracing::field::Field,
             }
             impl Visit for MessageVisitor<'_> {
+                fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
+                    if field == self.field {
+                        atrace::atrace_instant(self.tag, value);
+                    }
+                }
                 fn record_debug(
                     &mut self,
                     field: &tracing::field::Field,
                     value: &dyn std::fmt::Debug,
                 ) {
-                    if field.name() == "message" {
-                        self.s.push_str(&format!("{:?}", value));
+                    if field == self.field {
+                        atrace::atrace_instant(self.tag, &format!("{:?}", value));
                     }
                 }
             }
-            event.record(&mut MessageVisitor { s: &mut event_str });
+            event.record(&mut MessageVisitor { tag: self.tag, field: &field });
+        } else {
+            atrace::atrace_instant(
+                self.tag,
+                &format!("{} event", event.metadata().level().as_str()),
+            );
+        }
+    }
+}
+
+struct FieldFormatter {
+    is_event: bool,
+    s: String,
+}
+
+impl FieldFormatter {
+    fn new() -> FieldFormatter {
+        const DEFAULT_STR_CAPACITY: usize = 128; // Should fit most events without realloc.
+        FieldFormatter { is_event: true, s: String::with_capacity(DEFAULT_STR_CAPACITY) }
+    }
+
+    fn for_event() -> FieldFormatter {
+        FieldFormatter { is_event: true, ..FieldFormatter::new() }
+    }
+    fn for_span(span_name: &str) -> FieldFormatter {
+        let mut formatter = FieldFormatter { is_event: false, ..FieldFormatter::new() };
+        formatter.s.push_str(span_name);
+        formatter
+    }
+
+    fn as_str(&self) -> &str {
+        &self.s
+    }
+    fn add_delimeter_if_needed(&mut self) {
+        if !self.s.is_empty() {
+            self.s.push_str(", ");
         }
-        if event_str.is_empty() {
-            event_str = format!("{} event", event.metadata().level().as_str());
+    }
+}
+
+impl Visit for FieldFormatter {
+    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
+        self.add_delimeter_if_needed();
+        if self.is_event && field.name() == "message" {
+            self.s.push_str(value);
+        } else {
+            write!(&mut self.s, "{} = \"{}\"", field.name(), value).unwrap();
+        }
+    }
+    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
+        self.add_delimeter_if_needed();
+        if self.is_event && field.name() == "message" {
+            write!(&mut self.s, "{:?}", value).unwrap();
+        } else {
+            write!(&mut self.s, "{} = {:?}", field.name(), value).unwrap();
         }
-        atrace::atrace_instant(self.tag, &event_str);
     }
 }
 
@@ -165,6 +246,10 @@ mod tests {
         /// Implement this trait in the test with mocking logic and checks in implemented functions.
         /// Default implementations panic.
         pub trait ATraceMocker {
+            fn atrace_is_tag_enabled(&mut self, _tag: AtraceTag) -> bool {
+                panic!("Unexpected call");
+            }
+
             fn atrace_begin(&mut self, _tag: AtraceTag, _name: &str) {
                 panic!("Unexpected call");
             }
@@ -233,6 +318,9 @@ mod tests {
 
         // Wrapped functions that forward calls into mocker.
 
+        pub fn atrace_is_tag_enabled(tag: AtraceTag) -> bool {
+            with_mocker(|m| m.atrace_is_tag_enabled(tag))
+        }
         pub fn atrace_begin(tag: AtraceTag, name: &str) {
             with_mocker(|m| m.atrace_begin(tag, name))
         }
@@ -253,6 +341,9 @@ mod tests {
             begin_count: u32,
         }
         impl mock_atrace::ATraceMocker for CallCheck {
+            fn atrace_is_tag_enabled(&mut self, _tag: AtraceTag) -> bool {
+                true
+            }
             fn atrace_begin(&mut self, tag: AtraceTag, name: &str) {
                 self.begin_count += 1;
                 assert!(self.begin_count < 2);
@@ -283,6 +374,9 @@ mod tests {
             end_count: u32,
         }
         impl mock_atrace::ATraceMocker for CallCheck {
+            fn atrace_is_tag_enabled(&mut self, _tag: AtraceTag) -> bool {
+                true
+            }
             fn atrace_begin(&mut self, _tag: AtraceTag, _name: &str) {}
             fn atrace_end(&mut self, tag: AtraceTag) {
                 self.end_count += 1;
@@ -316,6 +410,9 @@ mod tests {
             end_count: u32,
         }
         impl mock_atrace::ATraceMocker for CallCheck {
+            fn atrace_is_tag_enabled(&mut self, _tag: AtraceTag) -> bool {
+                true
+            }
             fn atrace_begin(&mut self, _tag: AtraceTag, _name: &str) {
                 assert_eq!(self.end_count, 0);
                 assert_eq!(self.instant_count, 0);
@@ -365,6 +462,9 @@ mod tests {
             instant_count: u32,
         }
         impl mock_atrace::ATraceMocker for CallCheck {
+            fn atrace_is_tag_enabled(&mut self, _tag: AtraceTag) -> bool {
+                true
+            }
             fn atrace_instant(&mut self, tag: AtraceTag, name: &str) {
                 self.instant_count += 1;
                 assert!(self.instant_count < 2);
@@ -394,6 +494,9 @@ mod tests {
             instant_count: u32,
         }
         impl mock_atrace::ATraceMocker for CallCheck {
+            fn atrace_is_tag_enabled(&mut self, _tag: AtraceTag) -> bool {
+                true
+            }
             fn atrace_instant(&mut self, _tag: AtraceTag, name: &str) {
                 self.instant_count += 1;
                 assert!(self.instant_count < 2);
@@ -422,6 +525,9 @@ mod tests {
             instant_count: u32,
         }
         impl mock_atrace::ATraceMocker for CallCheck {
+            fn atrace_is_tag_enabled(&mut self, _tag: AtraceTag) -> bool {
+                true
+            }
             fn atrace_instant(&mut self, _tag: AtraceTag, name: &str) {
                 self.instant_count += 1;
                 assert!(self.instant_count < 2);
@@ -452,6 +558,9 @@ mod tests {
             end_count: u32,
         }
         impl mock_atrace::ATraceMocker for CallCheck {
+            fn atrace_is_tag_enabled(&mut self, _tag: AtraceTag) -> bool {
+                true
+            }
             fn atrace_begin(&mut self, tag: AtraceTag, _name: &str) {
                 self.begin_count += 1;
                 assert!(self.begin_count < 2);
@@ -496,6 +605,9 @@ mod tests {
             instant_count: u32,
         }
         impl mock_atrace::ATraceMocker for CallCheck {
+            fn atrace_is_tag_enabled(&mut self, _tag: AtraceTag) -> bool {
+                true
+            }
             fn atrace_begin(&mut self, _tag: AtraceTag, name: &str) {
                 self.begin_count += 1;
                 assert!(self.begin_count < 2);
@@ -531,6 +643,9 @@ mod tests {
             instant_count: u32,
         }
         impl mock_atrace::ATraceMocker for CallCheck {
+            fn atrace_is_tag_enabled(&mut self, _tag: AtraceTag) -> bool {
+                true
+            }
             fn atrace_instant(&mut self, _tag: AtraceTag, name: &str) {
                 self.instant_count += 1;
                 assert!(self.instant_count < 2);
@@ -558,6 +673,9 @@ mod tests {
             begin_count: u32,
         }
         impl mock_atrace::ATraceMocker for CallCheck {
+            fn atrace_is_tag_enabled(&mut self, _tag: AtraceTag) -> bool {
+                true
+            }
             fn atrace_begin(&mut self, _tag: AtraceTag, name: &str) {
                 self.begin_count += 1;
                 assert!(self.begin_count < 2);