From 81d1bb54930fb31f475370f6b0ab04978cba4e81 Mon Sep 17 00:00:00 2001
From: Jeremy Wall <jeremy@marzhillstudios.com>
Date: Sat, 19 Nov 2022 19:33:42 -0500
Subject: [PATCH] Correctly respect the instrumented spans

---
 Cargo.toml   |  5 ++++-
 src/lib.rs   | 38 +++++++++++++++++++++++++++++++++++---
 src/tests.rs | 30 ++++++++++++++++++++++++++----
 3 files changed, 65 insertions(+), 8 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 4cdec90..366e7bb 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -21,4 +21,7 @@ version = "0.3.1"
 features = ["registry"]
 
 [dependencies.wasm-bindgen]
-version = "0.2.81"
+version = "0.2.78"
+
+[dependencies.wasm-bindgen-test]
+version = "0.3"
diff --git a/src/lib.rs b/src/lib.rs
index a63a2f9..407a05e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -29,6 +29,7 @@ mod typings;
 pub struct BrowserLayer {
     record_timings: bool,
     max_level: tracing::Level,
+    //span_stack: Vec<span::Id>,
 }
 
 impl Default for BrowserLayer {
@@ -55,7 +56,13 @@ impl BrowserLayer {
         self
     }
 
-    pub fn write_for_level(&self, level: &tracing::Level, origin: String, recorder: Recorder) {
+    pub fn write_for_level(
+        &self,
+        level: &tracing::Level,
+        origin: String,
+        context: Option<Vec<String>>,
+        recorder: Recorder,
+    ) {
         let f: fn(String) = match *level {
             tracing::Level::ERROR => typings::error,
             tracing::Level::WARN => typings::warn,
@@ -63,7 +70,13 @@ impl BrowserLayer {
             tracing::Level::DEBUG => typings::debug,
             tracing::Level::TRACE => typings::trace,
         };
-        f(format!("{} {} {}", level, origin, recorder));
+        let spans = context
+            .map(|v| {
+                let s: String = v.concat();
+                s
+            })
+            .unwrap_or_else(|| String::new());
+        f(format!("{} {} {} {}", level, origin, recorder, spans));
     }
 }
 
@@ -97,16 +110,35 @@ impl<S: tracing::Subscriber + for<'a> LookupSpan<'a>> Layer<S> for BrowserLayer
         }
     }
 
-    fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) {
+    fn on_event(&self, event: &tracing::Event<'_>, ctx: Context<'_, S>) {
         let mut recorder = Recorder::new();
         event.record(&mut recorder);
         let metadata = event.metadata();
+        let recorders = ctx.current_span().id().and_then(|id| {
+            ctx.span_scope(id).map(|scope| {
+                let span_details: Vec<String> = scope
+                    .from_root()
+                    .map(|span_ref| {
+                        format!(
+                            "{} {}",
+                            span_ref.name(),
+                            span_ref
+                                .extensions()
+                                .get::<Recorder>()
+                                .expect("Unregistered Span")
+                        )
+                    })
+                    .collect();
+                span_details
+            })
+        });
         self.write_for_level(
             metadata.level(),
             metadata
                 .file()
                 .and_then(|file| metadata.line().map(|ln| format!("{}:{}", file, ln)))
                 .unwrap_or_default(),
+            recorders,
             recorder,
         );
     }
diff --git a/src/tests.rs b/src/tests.rs
index 69c2ce8..a4bb9c4 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -11,8 +11,30 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-#[test]
-fn it_works() {
-    let result = 2 + 2;
-    assert_eq!(result, 4);
+use tracing::{info, instrument};
+use wasm_bindgen_test::wasm_bindgen_test;
+
+use super::*;
+
+wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);
+
+#[instrument]
+fn instrumented_incrementor(arg: i32) -> i32 {
+    info!("Incrementing arg by 1");
+    arg + 1
+}
+
+#[instrument]
+fn instrumented_function(arg: i32) {
+    info!("Calling increment function");
+    let incremented = instrumented_incrementor(arg);
+    info!(incremented, "Finished incrementing function");
+}
+
+#[instrument]
+#[wasm_bindgen_test]
+fn it_works() {
+    configure_as_global_default();
+    info!("Calling instrumented function");
+    instrumented_function(3);
 }