From 211b9b02350cf85440a17737ba6d45e3813dbdca Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@buoyant.io>
Date: Fri, 17 Sep 2021 12:35:48 -0700
Subject: [PATCH] subscriber: fix missing `on_layer` for `Box` and `Arc` layer
 impls

The `Layer::on_layer` method on `Layer` was added in PR #1523. PR #1536,
which added `Layer` implementations to `Box<dyn Layer<...> + ...>` and
`Arc<dyn Layer<...> + ...>`, merged prior to #1523. However, when
merging #1523, I didn't think to update the `Layer` impl for `Box`ed and
`Arc`ed layers to forward `on_layer` to the inner `Layer`. This means
that when a `Layer` is wrapped in an `Arc` or a `Box`, the `on_layer`
method never gets called.

When per-layer filters are in use, the `on_layer` method is necessary to
ensure the filter is registered with the inner subscriber and has a
valid ID. This bug means that when per-layer filters are wrapped in a
`Box` or `Arc`, they won't be registered, and per-layer filtering
breaks.

This PR fixes the bug by adding `on_layer` implementations to the
`Layer` impls for `Arc`ed and `Box`ed layers. I also added some tests
--- thanks to @Noah-Kennedy for the original repro that these were based
on (https://github.com/tokio-rs/tracing/issues/1563#issuecomment-922014777).

Signed-off-by: Eliza Weisman <eliza@buoyant.io>
---
 tracing-subscriber/src/layer/mod.rs | 32 ++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/tracing-subscriber/src/layer/mod.rs b/tracing-subscriber/src/layer/mod.rs
index f3a87898e0..5ba89575b4 100644
--- a/tracing-subscriber/src/layer/mod.rs
+++ b/tracing-subscriber/src/layer/mod.rs
@@ -414,7 +414,11 @@
 //! [`LevelFilter`]: crate::filter::LevelFilter
 //! [feat]: crate#feature-flags
 use crate::filter;
-use std::{any::TypeId, ops::Deref, sync::Arc};
+use std::{
+    any::TypeId,
+    ops::{Deref, DerefMut},
+    sync::Arc,
+};
 use tracing_core::{
     metadata::Metadata,
     span,
@@ -1160,6 +1164,15 @@ where
     L: Layer<S>,
     S: Subscriber,
 {
+    fn on_layer(&mut self, subscriber: &mut S) {
+        if let Some(inner) = Arc::get_mut(self) {
+            // XXX(eliza): this may behave weird if another `Arc` clone of this
+            // layer is layered onto a _different_ subscriber...but there's no
+            // good solution for that...
+            inner.on_layer(subscriber);
+        }
+    }
+
     layer_impl_body! {}
 }
 
@@ -1167,6 +1180,15 @@ impl<S> Layer<S> for Arc<dyn Layer<S> + Send + Sync>
 where
     S: Subscriber,
 {
+    fn on_layer(&mut self, subscriber: &mut S) {
+        if let Some(inner) = Arc::get_mut(self) {
+            // XXX(eliza): this may behave weird if another `Arc` clone of this
+            // layer is layered onto a _different_ subscriber...but there's no
+            // good solution for that...
+            inner.on_layer(subscriber);
+        }
+    }
+
     layer_impl_body! {}
 }
 
@@ -1175,6 +1197,10 @@ where
     L: Layer<S>,
     S: Subscriber,
 {
+    fn on_layer(&mut self, subscriber: &mut S) {
+        self.deref_mut().on_layer(subscriber);
+    }
+
     layer_impl_body! {}
 }
 
@@ -1182,6 +1208,10 @@ impl<S> Layer<S> for Box<dyn Layer<S> + Send + Sync>
 where
     S: Subscriber,
 {
+    fn on_layer(&mut self, subscriber: &mut S) {
+        self.deref_mut().on_layer(subscriber);
+    }
+
     layer_impl_body! {}
 }