subscriber: fix span data for new, exit, and close events (#1334)

* subscriber: fix span data for new, exit, and close events

New, exit and close span events are generated while the current
context is set to either `None` or the parent span of the span the
event relates to. This causes spans data to be absent from the JSON
output in the case of the `None`, or causes the span data to reference
the parent's span data. Changing the way the current span is
determined allows the correct span to be identified for these
events. Trying to access the events `.parent()` allows access of the
correct span for the `on_event` actions, while using `.current_span()`
works for normal events.

Ref: #1032

* Fix style

* subscriber: improve test for #1333

Based on feedback by @hawkw, I've improved the test for #1333 to parse
the json output. This is more specifc for the bug and allows easier
testing of the different span `on_events`.

Ref: #1333 (review)

* subscriber: improve #1334 tests covering all span states

Use the `on_records` test method check all events have the correct
context as described in the PR.

* Apply suggestions from code review

Co-authored-by: Eliza Weisman <eliza@buoyant.io>

tracing-subscriber/src/fmt/format/json.rs

@@ -209,7 +209,10 @@ where
 
             let current_span = if self.format.display_current_span || self.format.display_span_list
             {
-                ctx.ctx.current_span().id().and_then(|id| ctx.ctx.span(id))
+                event
+                    .parent()
+                    .and_then(|id| ctx.span(id))
+                    .or_else(|| ctx.lookup_current())
             } else {
                 None
             };
@@ -497,6 +500,7 @@ impl<'a> fmt::Debug for WriteAdaptor<'a> {
 mod test {
     use super::*;
     use crate::fmt::{format::FmtSpan, test::MockMakeWriter, time::FormatTime, CollectorBuilder};
+
     use tracing::{self, collect::with_default};
 
     use std::fmt;
@@ -613,31 +617,13 @@ mod test {
         // This test reproduces issue #707, where using `Span::record` causes
         // any events inside the span to be ignored.
 
-        let make_writer = MockMakeWriter::default();
-        let subscriber = crate::fmt()
-            .json()
-            .with_writer(make_writer.clone())
-            .finish();
-
-        let parse_buf = || -> serde_json::Value {
-            let buf = String::from_utf8(make_writer.buf().to_vec()).unwrap();
-            let json = buf
-                .lines()
-                .last()
-                .expect("expected at least one line to be written!");
-            match serde_json::from_str(&json) {
-                Ok(v) => v,
-                Err(e) => panic!(
-                    "assertion failed: JSON shouldn't be malformed\n  error: {}\n  json: {}",
-                    e, json
-                ),
-            }
-        };
+        let buffer = MockMakeWriter::default();
+        let subscriber = crate::fmt().json().with_writer(buffer.clone()).finish();
 
         with_default(subscriber, || {
             tracing::info!("an event outside the root span");
             assert_eq!(
-                parse_buf()["fields"]["message"],
+                parse_as_json(&buffer)["fields"]["message"],
                 "an event outside the root span"
             );
 
@@ -647,29 +633,113 @@ mod test {
 
             tracing::info!("an event inside the root span");
             assert_eq!(
-                parse_buf()["fields"]["message"],
+                parse_as_json(&buffer)["fields"]["message"],
                 "an event inside the root span"
             );
         });
     }
 
     #[test]
-    fn json_span_event() {
+    fn json_span_event_show_correct_context() {
+        let buffer = MockMakeWriter::default();
+        let subscriber = collector()
+            .with_writer(buffer.clone())
+            .flatten_event(false)
+            .with_current_span(true)
+            .with_span_list(false)
+            .with_span_events(FmtSpan::FULL)
+            .finish();
+
+        with_default(subscriber, || {
+            let context = "parent";
+            let parent_span = tracing::info_span!("parent_span", context);
+
+            let event = parse_as_json(&buffer);
+            assert_eq!(event["fields"]["message"], "new");
+            assert_eq!(event["span"]["context"], "parent");
+
+            let _parent_enter = parent_span.enter();
+            let event = parse_as_json(&buffer);
+            assert_eq!(event["fields"]["message"], "enter");
+            assert_eq!(event["span"]["context"], "parent");
+
+            let context = "child";
+            let child_span = tracing::info_span!("child_span", context);
+            let event = parse_as_json(&buffer);
+            assert_eq!(event["fields"]["message"], "new");
+            assert_eq!(event["span"]["context"], "child");
+
+            let _child_enter = child_span.enter();
+            let event = parse_as_json(&buffer);
+            assert_eq!(event["fields"]["message"], "enter");
+            assert_eq!(event["span"]["context"], "child");
+
+            drop(_child_enter);
+            let event = parse_as_json(&buffer);
+            assert_eq!(event["fields"]["message"], "exit");
+            assert_eq!(event["span"]["context"], "child");
+
+            drop(child_span);
+            let event = parse_as_json(&buffer);
+            assert_eq!(event["fields"]["message"], "close");
+            assert_eq!(event["span"]["context"], "child");
+
+            drop(_parent_enter);
+            let event = parse_as_json(&buffer);
+            assert_eq!(event["fields"]["message"], "exit");
+            assert_eq!(event["span"]["context"], "parent");
+
+            drop(parent_span);
+            let event = parse_as_json(&buffer);
+            assert_eq!(event["fields"]["message"], "close");
+            assert_eq!(event["span"]["context"], "parent");
+        });
+    }
+
+    #[test]
+    fn json_span_event_with_no_fields() {
         // Check span events serialize correctly.
         // Discussion: /~https://github.com/tokio-rs/tracing/issues/829#issuecomment-661984255
         //
-        let expected = r#"{"timestamp":"fake time","level":"INFO","fields":{"message":"enter"},"target":"tracing_subscriber::fmt::format::json::test"}"#;
-        let collector = collector()
+        let buffer = MockMakeWriter::default();
+        let subscriber = collector()
+            .with_writer(buffer.clone())
             .flatten_event(false)
             .with_current_span(false)
             .with_span_list(false)
-            .with_span_events(FmtSpan::ENTER);
+            .with_span_events(FmtSpan::FULL)
+            .finish();
 
-        test_json(expected, collector, || {
-            tracing::info_span!("valid_json").in_scope(|| {});
+        with_default(subscriber, || {
+            let span = tracing::info_span!("valid_json");
+            assert_eq!(parse_as_json(&buffer)["fields"]["message"], "new");
+
+            let _enter = span.enter();
+            assert_eq!(parse_as_json(&buffer)["fields"]["message"], "enter");
+
+            drop(_enter);
+            assert_eq!(parse_as_json(&buffer)["fields"]["message"], "exit");
+
+            drop(span);
+            assert_eq!(parse_as_json(&buffer)["fields"]["message"], "close");
         });
     }
 
+    fn parse_as_json(buffer: &MockMakeWriter) -> serde_json::Value {
+        let buf = String::from_utf8(buffer.buf().to_vec()).unwrap();
+        let json = buf
+            .lines()
+            .last()
+            .expect("expected at least one line to be written!");
+        match serde_json::from_str(&json) {
+            Ok(v) => v,
+            Err(e) => panic!(
+                "assertion failed: JSON shouldn't be malformed\n  error: {}\n  json: {}",
+                e, json
+            ),
+        }
+    }
+
     fn test_json<T>(
         expected: &str,
         builder: crate::fmt::CollectorBuilder<JsonFields, Format<Json>>,