rust: Cache span names as `CString`s

We also avoid parsing span names as `CString`s if tracing is disabled.

sdk-lib/src/main/rust/utils/trace.rs

@@ -1,4 +1,8 @@
-use std::ffi::CString;
+use std::{
+    collections::HashMap,
+    ffi::CString,
+    sync::{Arc, Mutex},
+};
 
 use tracing::{error, span, Subscriber};
 use tracing_subscriber::{layer::Context, registry::LookupSpan};
@@ -7,16 +11,53 @@ use super::target_ndk::{Api23, NdkApi};
 
 pub struct Layer {
     ndk_api: Option<NdkApi>,
+    open_spans: Arc<Mutex<HashMap<span::Id, CString>>>,
 }
 
 impl Layer {
     pub fn new(ndk_api: Option<NdkApi>) -> Self {
-        Layer { ndk_api }
+        Layer {
+            ndk_api,
+            open_spans: Arc::new(Mutex::new(HashMap::new())),
+        }
     }
 
-    fn with_api(&self, f: impl FnOnce(&Api23)) {
-        if let Some(v23) = self.ndk_api.as_ref().and_then(|api| api.v23.as_ref()) {
-            f(v23)
+    fn with_api(&self, f23: impl FnOnce(&Api23)) {
+        if let Some(api) = self.ndk_api.as_ref() {
+            if let Some(v23) = api.v23.as_ref() {
+                if unsafe { v23.ATrace_isEnabled() } {
+                    f23(v23)
+                }
+            }
+        }
+    }
+
+    fn with_entered_span<S: Subscriber>(
+        &self,
+        id: &span::Id,
+        ctx: &Context<'_, S>,
+        f: impl FnOnce(&CString),
+    ) where
+        for<'lookup> S: LookupSpan<'lookup>,
+    {
+        let mut open_spans = self.open_spans.lock().unwrap();
+
+        if let Some(section_name) = open_spans.get_mut(id) {
+            f(&section_name);
+        } else {
+            // We need to obtain the span's name as a CString.
+            match ctx.metadata(id) {
+                Some(metadata) => match CString::new(metadata.name()) {
+                    Ok(section_name) => {
+                        f(&section_name);
+                    }
+                    Err(_) => error!(
+                        "Span name contains internal NUL byte: '{}'",
+                        metadata.name()
+                    ),
+                },
+                None => error!("Span {:?} has no metadata", id),
+            }
         }
     }
 }
@@ -26,15 +67,10 @@ where
     for<'lookup> S: LookupSpan<'lookup>,
 {
     fn on_enter(&self, id: &span::Id, ctx: Context<'_, S>) {
-        self.with_api(|api| match ctx.metadata(id) {
-            Some(metadata) => match CString::new(metadata.name()) {
-                Ok(section_name) => unsafe { api.ATrace_beginSection(section_name.as_ptr()) },
-                Err(_) => error!(
-                    "Span name contains internal NUL byte: '{}'",
-                    metadata.name()
-                ),
-            },
-            None => error!("Span {:?} has no metadata", id),
+        self.with_api(|api| {
+            self.with_entered_span(id, &ctx, |section_name| unsafe {
+                api.ATrace_beginSection(section_name.as_ptr())
+            })
         });
     }
 
@@ -43,4 +79,11 @@ where
             unsafe { api.ATrace_endSection() };
         });
     }
+
+    fn on_id_change(&self, old: &span::Id, new: &span::Id, _ctx: Context<'_, S>) {
+        let mut open_spans = self.open_spans.lock().unwrap();
+        if let Some(value) = open_spans.remove(old) {
+            open_spans.insert(new.clone(), value);
+        }
+    }
 }