From a8489adf1c2c531fc61c5d6a1eca40e07e79a3f8 Mon Sep 17 00:00:00 2001
From: David Barsky <me@davidbarsky.com>
Date: Sat, 27 Nov 2021 17:15:22 -0500
Subject: [PATCH] subscriber: forward `Filtered::downcast_raw` to wrapped Layer
 (#1619)

## Motivation

I'm trying to implement a `Layer` that has something very similar to
tracing-error's
[`WithContext`](https://github.com/tokio-rs/tracing/blob/66cd79f72af5ebcb6f21a1017b6ce33bea05558d/tracing-error/src/layer.rs#L32)
to support erasing types and getting access to the current span's state.

To do that, I implement `Layer::downcast_raw` in [the same way that
tracing-error
does](https://github.com/tokio-rs/tracing/blob/66cd79f72af5ebcb6f21a1017b6ce33bea05558d/tracing-error/src/layer.rs#L55-L63).

This works great when the layer is not filtered. However, once I filter
the layer

```rust
let filter = tracing_subscriber::filter::Targets::new().with_default(tracing::Level::INFO);
let layer = MyLayer::new();
tracing_subscriber::registry().with(layer.with_filter(filter)).init();
```

I'm not able to get a `WithContext` instance anymore, because `Filtered`
[handles `downcast_raw`, and doesn't forward
it](https://github.com/tokio-rs/tracing/blob/66cd79f72af5ebcb6f21a1017b6ce33bea05558d/tracing-subscriber/src/filter/layer_filters.rs#L379-L391)
to `MyLayer::downcast_raw`.

## Solution

If `Filtered::downcast_raw` does not know how to handle the given type,
forward it to the wrapped layer's `Layer::downcast_raw` implementation.

Fixes #1618
---
 .../src/filter/layer_filters.rs               |  2 +-
 .../tests/layer_filters/downcast_raw.rs       | 68 +++++++++++++++++++
 .../tests/layer_filters/main.rs               |  1 +
 3 files changed, 70 insertions(+), 1 deletion(-)
 create mode 100644 tracing-subscriber/tests/layer_filters/downcast_raw.rs

diff --git a/tracing-subscriber/src/filter/layer_filters.rs b/tracing-subscriber/src/filter/layer_filters.rs
index 5986dee14e..e0ad797585 100644
--- a/tracing-subscriber/src/filter/layer_filters.rs
+++ b/tracing-subscriber/src/filter/layer_filters.rs
@@ -380,7 +380,7 @@ where
             id if id == TypeId::of::<MagicPlfDowncastMarker>() => {
                 Some(&self.id as *const _ as *const ())
             }
-            _ => None,
+            _ => self.layer.downcast_raw(id),
         }
     }
 }
diff --git a/tracing-subscriber/tests/layer_filters/downcast_raw.rs b/tracing-subscriber/tests/layer_filters/downcast_raw.rs
new file mode 100644
index 0000000000..b5f7e35ced
--- /dev/null
+++ b/tracing-subscriber/tests/layer_filters/downcast_raw.rs
@@ -0,0 +1,68 @@
+use tracing::Subscriber;
+use tracing_subscriber::filter::Targets;
+use tracing_subscriber::prelude::*;
+use tracing_subscriber::Layer;
+
+#[test]
+fn downcast_ref_to_inner_layer_and_filter() {
+    // Test that a filtered layer gives downcast_ref access to
+    // both the layer and the filter.
+
+    struct WrappedLayer;
+
+    impl<S> Layer<S> for WrappedLayer
+    where
+        S: Subscriber,
+        S: for<'lookup> tracing_subscriber::registry::LookupSpan<'lookup>,
+    {
+    }
+
+    let layer = WrappedLayer;
+    let filter = Targets::new().with_default(tracing::Level::INFO);
+    let registry = tracing_subscriber::registry().with(layer.with_filter(filter));
+    let dispatch = tracing::dispatcher::Dispatch::new(registry);
+
+    // The filter is available
+    assert!(dispatch.downcast_ref::<Targets>().is_some());
+    // The wrapped layer is available
+    assert!(dispatch.downcast_ref::<WrappedLayer>().is_some());
+}
+
+#[test]
+fn forward_downcast_raw_to_layer() {
+    // Test that a filtered layer still gives its wrapped layer a chance to
+    // return a custom struct from downcast_raw.
+    // https://github.com/tokio-rs/tracing/issues/1618
+
+    struct WrappedLayer {
+        with_context: WithContext,
+    }
+
+    struct WithContext;
+
+    impl<S> Layer<S> for WrappedLayer
+    where
+        S: Subscriber,
+        S: for<'lookup> tracing_subscriber::registry::LookupSpan<'lookup>,
+    {
+        unsafe fn downcast_raw(&self, id: std::any::TypeId) -> Option<*const ()> {
+            match id {
+                id if id == std::any::TypeId::of::<Self>() => Some(self as *const _ as *const ()),
+                id if id == std::any::TypeId::of::<WithContext>() => {
+                    Some(&self.with_context as *const _ as *const ())
+                }
+                _ => None,
+            }
+        }
+    }
+
+    let layer = WrappedLayer {
+        with_context: WithContext,
+    };
+    let filter = Targets::new().with_default(tracing::Level::INFO);
+    let registry = tracing_subscriber::registry().with(layer.with_filter(filter));
+    let dispatch = tracing::dispatcher::Dispatch::new(registry);
+
+    // Types from a custom implementation of `downcast_raw` are available
+    assert!(dispatch.downcast_ref::<WithContext>().is_some());
+}
diff --git a/tracing-subscriber/tests/layer_filters/main.rs b/tracing-subscriber/tests/layer_filters/main.rs
index ccc4ef2679..348c3fec91 100644
--- a/tracing-subscriber/tests/layer_filters/main.rs
+++ b/tracing-subscriber/tests/layer_filters/main.rs
@@ -2,6 +2,7 @@
 #[path = "../support.rs"]
 mod support;
 use self::support::*;
+mod downcast_raw;
 mod filter_scopes;
 mod targets;
 mod trees;