Don't call wrapped atrace_begin if tag is disabled

Gating the binding call on tag status lets us avoid constructing CString
and spending a lot of time when tracing is disabled.

We'll use this optimization when adding other wrapper methods that
accept strings.

Change of atrace_begin time with tracing off: 100 ns -> 6 ns

Test: atest
Test: Manually run libatrace_rust_benchmark
Test: Manually run libatrace_rust_example and collect a trace with
      Perfetto.

Bug: 289989828
Change-Id: Id73cdfa887da9b5d2142154d3be5f88872108675

2 files changed, 99 insertions, 18 deletions

diff --git a/libatrace_rust/benchmark/README.md b/libatrace_rust/benchmark/README.md
index c4e35fe0..27af5176 100644
--- a/libatrace_rust/benchmark/README.md
+++ b/libatrace_rust/benchmark/README.md
@@ -36,12 +36,12 @@ device may be different but we expect similar relative performance between Rust
 Rust results from `libatrace_rust_benchmark 2>&1 | grep time`:
 
 ```text
-tracing_off_begin/10    time:   [91.195 ns 91.648 ns 92.193 ns]
-tracing_off_begin/1000  time:   [289.50 ns 290.57 ns 291.76 ns]
-tracing_off_end         time:   [6.0455 ns 6.0998 ns 6.1644 ns]
-tracing_on_begin/10     time:   [1.2285 µs 1.2330 µs 1.2379 µs]
-tracing_on_begin/1000   time:   [1.4667 µs 1.4709 µs 1.4754 µs]
-tracing_on_end          time:   [1.1344 µs 1.1444 µs 1.1543 µs]
+tracing_off_begin/10    time:   [6.0211 ns 6.0382 ns 6.0607 ns]
+tracing_off_begin/1000  time:   [6.0119 ns 6.0418 ns 6.0823 ns]
+tracing_off_end         time:   [6.5417 ns 6.6801 ns 6.8131 ns]
+tracing_on_begin/10     time:   [1.2847 µs 1.2929 µs 1.3044 µs]
+tracing_on_begin/1000   time:   [1.5395 µs 1.5476 µs 1.5580 µs]
+tracing_on_end          time:   [1.1153 µs 1.1208 µs 1.1276 µs]
 ```
 
 C++ results from `libatrace_rust_benchmark_cc`:
diff --git a/libatrace_rust/src/lib.rs b/libatrace_rust/src/lib.rs
index 102ec51c..8ecec7f8 100644
--- a/libatrace_rust/src/lib.rs
+++ b/libatrace_rust/src/lib.rs
@@ -85,10 +85,22 @@ pub mod tags {
     const_assert_eq!(AtraceTag::Thermal.bits(), cutils_trace_bindgen::ATRACE_TAG_LAST as u64);
 }
 
+/// Test if a given tag is currently enabled.
+///
+/// It can be used as a guard condition around more expensive trace calculations.
+pub fn atrace_is_tag_enabled(tag: AtraceTag) -> bool {
+    // SAFETY: No pointers are transferred.
+    unsafe { trace_bind::atrace_is_tag_enabled_wrap(tag.bits()) != 0 }
+}
+
 /// Trace the beginning of a context. `name` is used to identify the context.
 ///
 /// This is often used to time function execution.
 pub fn atrace_begin(tag: AtraceTag, name: &str) {
+    if !atrace_is_tag_enabled(tag) {
+        return;
+    }
+
     let name_cstr = CString::new(name.as_bytes()).expect("CString::new failed");
     // SAFETY: The function does not accept the pointer ownership, only reads its contents.
     // The passed string is guaranteed to be null-terminated by CString.
@@ -130,6 +142,9 @@ 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_wrap(&mut self, _tag: u64) -> u64 {
+                panic!("Unexpected call");
+            }
             fn atrace_begin_wrap(&mut self, _tag: u64, _name: *const c_char) {
                 panic!("Unexpected call");
             }
@@ -195,6 +210,9 @@ mod tests {
         // The functions are marked as unsafe to match the binding interface, won't compile otherwise.
         // The mocker methods themselves are not marked as unsafe.
 
+        pub unsafe fn atrace_is_tag_enabled_wrap(tag: u64) -> u64 {
+            with_mocker(|m| m.atrace_is_tag_enabled_wrap(tag))
+        }
         pub unsafe fn atrace_begin_wrap(tag: u64, name: *const c_char) {
             with_mocker(|m| m.atrace_begin_wrap(tag, name))
         }
@@ -207,13 +225,16 @@ mod tests {
     fn forwards_trace_begin() {
         #[derive(Default)]
         struct CallCheck {
-            called_count: u32,
+            begin_count: u32,
         }
 
         impl mock_atrace::ATraceMocker for CallCheck {
+            fn atrace_is_tag_enabled_wrap(&mut self, _tag: u64) -> u64 {
+                1
+            }
             fn atrace_begin_wrap(&mut self, tag: u64, name: *const c_char) {
-                self.called_count += 1;
-                assert!(self.called_count < 2);
+                self.begin_count += 1;
+                assert!(self.begin_count < 2);
                 assert_eq!(tag, cutils_trace_bindgen::ATRACE_TAG_APP as u64);
                 // SAFETY: If the code under test is correct, the pointer is guaranteed to satisfy
                 // the requirements of `CStr::from_ptr`. If the code is not correct, this section is
@@ -224,7 +245,7 @@ mod tests {
             }
 
             fn finish(&self) {
-                assert_eq!(self.called_count, 1);
+                assert_eq!(self.begin_count, 1);
             }
         }
 
@@ -236,21 +257,50 @@ mod tests {
     }
 
     #[test]
+    fn trace_begin_not_called_with_disabled_tag() {
+        #[derive(Default)]
+        struct CallCheck {
+            is_tag_enabled_count: u32,
+        }
+
+        impl mock_atrace::ATraceMocker for CallCheck {
+            fn atrace_is_tag_enabled_wrap(&mut self, _tag: u64) -> u64 {
+                self.is_tag_enabled_count += 1;
+                assert!(self.is_tag_enabled_count < 2);
+                0
+            }
+            fn atrace_begin_wrap(&mut self, _tag: u64, _name: *const c_char) {
+                panic!("Begin should not be called with disabled tag.")
+            }
+
+            fn finish(&self) {
+                assert_eq!(self.is_tag_enabled_count, 1);
+            }
+        }
+
+        let _guard = mock_atrace::set_scoped_mocker(CallCheck::default());
+
+        atrace_begin(AtraceTag::App, "Ignore me");
+
+        mock_atrace::mocker_finish();
+    }
+
+    #[test]
     fn forwards_trace_end() {
         #[derive(Default)]
         struct CallCheck {
-            called_count: u32,
+            end_count: u32,
         }
 
         impl mock_atrace::ATraceMocker for CallCheck {
             fn atrace_end_wrap(&mut self, tag: u64) {
-                self.called_count += 1;
-                assert!(self.called_count < 2);
+                self.end_count += 1;
+                assert!(self.end_count < 2);
                 assert_eq!(tag, cutils_trace_bindgen::ATRACE_TAG_APP as u64);
             }
 
             fn finish(&self) {
-                assert_eq!(self.called_count, 1);
+                assert_eq!(self.end_count, 1);
             }
         }
 
@@ -265,13 +315,16 @@ mod tests {
     fn can_combine_tags() {
         #[derive(Default)]
         struct CallCheck {
-            called_count: u32,
+            begin_count: u32,
         }
 
         impl mock_atrace::ATraceMocker for CallCheck {
+            fn atrace_is_tag_enabled_wrap(&mut self, _tag: u64) -> u64 {
+                1
+            }
             fn atrace_begin_wrap(&mut self, tag: u64, _name: *const c_char) {
-                self.called_count += 1;
-                assert!(self.called_count < 2);
+                self.begin_count += 1;
+                assert!(self.begin_count < 2);
                 assert_eq!(
                     tag,
                     (cutils_trace_bindgen::ATRACE_TAG_HAL | cutils_trace_bindgen::ATRACE_TAG_CAMERA)
@@ -280,7 +333,7 @@ mod tests {
             }
 
             fn finish(&self) {
-                assert_eq!(self.called_count, 1);
+                assert_eq!(self.begin_count, 1);
             }
         }
 
@@ -290,4 +343,32 @@ mod tests {
 
         mock_atrace::mocker_finish();
     }
+
+    #[test]
+    fn forwards_is_tag_enabled() {
+        #[derive(Default)]
+        struct CallCheck {
+            is_tag_enabled_count: u32,
+        }
+
+        impl mock_atrace::ATraceMocker for CallCheck {
+            fn atrace_is_tag_enabled_wrap(&mut self, tag: u64) -> u64 {
+                self.is_tag_enabled_count += 1;
+                assert!(self.is_tag_enabled_count < 2);
+                assert_eq!(tag, cutils_trace_bindgen::ATRACE_TAG_ADB as u64);
+                1
+            }
+
+            fn finish(&self) {
+                assert_eq!(self.is_tag_enabled_count, 1);
+            }
+        }
+
+        let _guard = mock_atrace::set_scoped_mocker(CallCheck::default());
+
+        let res = atrace_is_tag_enabled(AtraceTag::Adb);
+        assert!(res);
+
+        mock_atrace::mocker_finish();
+    }
 }