fmt: Track duplicate spans in fmt context (#361)

## Motivation

Entering twice in the same span causes the span to be lost after the
second span got droped, see #358 

## Solution

Duplicate spans are tracked in the `CONTEXT` thread local stack  and
ignored when finding the current span id. Spans are always pushed on the
`CONTEXT` stack. This way, when a span is dropped, its ref count
reflects the content of the thread local

Fixes #358

tracing-subscriber/src/fmt/span.rs

@@ -7,6 +7,7 @@ use std::{
 use crate::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 use owning_ref::OwningHandle;
 
+use std::collections::HashSet;
 pub(crate) use tracing_core::span::{Attributes, Current, Id, Record};
 use tracing_core::{dispatcher, Metadata};
 
@@ -62,8 +63,56 @@ enum State {
     Empty(usize),
 }
 
+struct ContextId {
+    id: Id,
+    duplicate: bool,
+}
+
+struct SpanStack {
+    stack: Vec<ContextId>,
+    ids: HashSet<Id>,
+}
+
+impl SpanStack {
+    fn new() -> Self {
+        SpanStack {
+            stack: vec![],
+            ids: HashSet::new(),
+        }
+    }
+
+    fn push(&mut self, id: Id) {
+        let duplicate = self.ids.contains(&id);
+        if !duplicate {
+            self.ids.insert(id.clone());
+        }
+        self.stack.push(ContextId { id, duplicate })
+    }
+
+    fn pop(&mut self, expected_id: &Id) -> Option<Id> {
+        if &self.stack.last()?.id == expected_id {
+            let ContextId { id, duplicate } = self.stack.pop()?;
+            if !duplicate {
+                self.ids.remove(&id);
+            }
+            Some(id)
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    fn current(&self) -> Option<&Id> {
+        self.stack
+            .iter()
+            .rev()
+            .find(|context_id| !context_id.duplicate)
+            .map(|context_id| &context_id.id)
+    }
+}
+
 thread_local! {
-    static CONTEXT: RefCell<Vec<Id>> = RefCell::new(vec![]);
+    static CONTEXT: RefCell<SpanStack> = RefCell::new(SpanStack::new());
 }
 
 macro_rules! debug_panic {
@@ -158,7 +207,7 @@ impl<'a, N> Context<'a, N> {
     {
         CONTEXT
             .try_with(|current| {
-                if let Some(id) = current.borrow().last() {
+                if let Some(id) = current.borrow().current() {
                     if let Some(span) = self.store.get(id) {
                         // with_parent uses the call stack to visit the span
                         // stack in reverse order, without having to allocate
@@ -183,9 +232,9 @@ impl<'a, N> Context<'a, N> {
         // will just do nothing rather than cause a double panic.
         CONTEXT
             .try_with(|current| {
-                if let Some(id) = current.borrow().last() {
-                    if let Some(span) = self.store.get(id) {
-                        return Some(f((id, span)));
+                if let Some(id) = current.borrow().current() {
+                    if let Some(span) = self.store.get(&id) {
+                        return Some(f((&id, span)));
                     } else {
                         debug_panic!("missing span for {:?}, this is a bug", id);
                     }
@@ -236,31 +285,17 @@ impl Store {
     #[inline]
     pub(crate) fn current(&self) -> Option<Id> {
         CONTEXT
-            .try_with(|current| current.borrow().last().map(|span| self.clone_span(span)))
+            .try_with(|current| current.borrow().current().map(|id| self.clone_span(id)))
             .ok()?
     }
 
     pub(crate) fn push(&self, id: &Id) {
-        let _ = CONTEXT.try_with(|current| {
-            let mut current = current.borrow_mut();
-            if current.contains(id) {
-                // Ignore duplicate enters.
-                return;
-            }
-            current.push(self.clone_span(id));
-        });
+        let _ = CONTEXT.try_with(|current| current.borrow_mut().push(self.clone_span(id)));
     }
 
     pub(crate) fn pop(&self, expected_id: &Id) {
         let id = CONTEXT
-            .try_with(|current| {
-                let mut current = current.borrow_mut();
-                if current.last() == Some(expected_id) {
-                    current.pop()
-                } else {
-                    None
-                }
-            })
+            .try_with(|current| current.borrow_mut().pop(expected_id))
             .ok()
             .and_then(|i| i);
         if let Some(id) = id {

tracing-subscriber/tests/duplicate_spans.rs

@@ -0,0 +1,47 @@
+mod support;
+use tracing::{self, subscriber::with_default, Span};
+use tracing_subscriber::{filter::EnvFilter, FmtSubscriber};
+
+#[test]
+fn duplicate_spans() {
+    let subscriber = FmtSubscriber::builder()
+        .with_env_filter(EnvFilter::new("[root]=debug"))
+        .finish();
+
+    with_default(subscriber, || {
+        let root = tracing::debug_span!("root");
+        root.in_scope(|| {
+            // root:
+            assert_eq!(root, Span::current(), "Current span must be 'root'");
+            let leaf = tracing::debug_span!("leaf");
+            leaf.in_scope(|| {
+                // root:leaf:
+                assert_eq!(leaf, Span::current(), "Current span must be 'leaf'");
+                root.in_scope(|| {
+                    // root:leaf:
+                    assert_eq!(
+                        leaf,
+                        Span::current(),
+                        "Current span must be 'leaf' after entering twice the 'root' span"
+                    );
+                })
+            });
+            // root:
+            assert_eq!(
+                root,
+                Span::current(),
+                "Current span must be root ('leaf' exited, nested 'root' exited)"
+            );
+
+            root.in_scope(|| {
+                assert_eq!(root, Span::current(), "Current span must be root");
+            });
+            // root:
+            assert_eq!(
+                root,
+                Span::current(),
+                "Current span must still be root after exiting nested 'root'"
+            );
+        });
+    });
+}