Fix tracing spans are not being forwarded to spawned task (#8009)

* Fix tracing spans are not being forwarded to spawned task

There is a bug that tracing spans are not forwarded to spawned task. The
problem was that only the telemetry span was forwarded. The solution to
this is to use the tracing provided `in_current_span` to capture the
current active span and pass the telemetry span explictely. We will now
always enter the span when the future is polled. This is essentially the
same strategy as tracing is doing with its `Instrumented`, but now
extended for our use case with having multiple spans active.

* More tests

substrate/client/service/src/task_manager/tests.rs

@@ -20,9 +20,10 @@ use crate::config::TaskExecutor;
 use crate::task_manager::TaskManager;
 use futures::{future::FutureExt, pin_mut, select};
 use parking_lot::Mutex;
-use std::any::Any;
-use std::sync::Arc;
-use std::time::Duration;
+use std::{any::Any, sync::Arc, time::Duration};
+use tracing_subscriber::{layer::{SubscriberExt, Context}, Layer};
+use tracing::{subscriber::Subscriber, span::{Attributes, Id, Record, Span}, event::Event};
+use sc_telemetry::TelemetrySpan;
 
 #[derive(Clone, Debug)]
 struct DropTester(Arc<Mutex<usize>>);
@@ -312,3 +313,94 @@ fn ensure_task_manager_future_continues_when_childs_not_essential_task_fails() {
 	runtime.block_on(task_manager.clean_shutdown());
 	assert_eq!(drop_tester, 0);
 }
+
+struct TestLayer {
+	spans_entered: Arc<Mutex<Vec<String>>>,
+	spans: Arc<Mutex<std::collections::HashMap<Id, String>>>,
+}
+
+impl<S: Subscriber> Layer<S> for TestLayer {
+	fn new_span(&self, attrs: &Attributes<'_>, id: &Id, _ctx: Context<S>) {
+		self.spans.lock().insert(id.clone(), attrs.metadata().name().to_string());
+	}
+
+	fn on_record(&self, _: &Id, _: &Record<'_>, _: Context<S>) {}
+
+	fn on_event(&self, _: &Event<'_>, _: Context<S>) {}
+
+	fn on_enter(&self, span: &Id, _: Context<S>) {
+		let name = self.spans.lock().get(span).unwrap().clone();
+		self.spans_entered.lock().push(name);
+	}
+
+	fn on_exit(&self, _: &Id, _: Context<S>) {}
+
+	fn on_close(&self, _: Id, _: Context<S>) {}
+}
+
+type TestSubscriber = tracing_subscriber::layer::Layered<
+	TestLayer,
+	tracing_subscriber::fmt::Subscriber
+>;
+
+fn setup_subscriber() -> (
+	TestSubscriber,
+	Arc<Mutex<Vec<String>>>,
+) {
+	let spans_entered = Arc::new(Mutex::new(Default::default()));
+	let layer = TestLayer {
+		spans: Arc::new(Mutex::new(Default::default())),
+		spans_entered: spans_entered.clone(),
+	};
+	let subscriber = tracing_subscriber::fmt().finish().with(layer);
+	(subscriber, spans_entered)
+}
+
+#[test]
+fn telemetry_span_is_forwarded_to_task() {
+	let (subscriber, spans_entered) = setup_subscriber();
+	let _sub_guard = tracing::subscriber::set_global_default(subscriber);
+
+	let telemetry_span = TelemetrySpan::new();
+
+	let span = tracing::info_span!("test");
+	let _enter = span.enter();
+
+	let mut runtime = tokio::runtime::Runtime::new().unwrap();
+	let handle = runtime.handle().clone();
+	let task_executor = TaskExecutor::from(move |fut, _| handle.spawn(fut).map(|_| ()));
+	let task_manager = TaskManager::new(task_executor, None, Some(telemetry_span.clone())).unwrap();
+
+	let (sender, receiver) = futures::channel::oneshot::channel();
+	let spawn_handle = task_manager.spawn_handle();
+
+	let span = span.clone();
+	task_manager.spawn_handle().spawn(
+		"test",
+		async move {
+			assert_eq!(span, Span::current());
+			spawn_handle.spawn("test-nested", async move {
+				assert_eq!(span, Span::current());
+				sender.send(()).unwrap();
+			}.boxed());
+		}.boxed(),
+	);
+
+	// We need to leave exit the span here. If tokio is not running with multithreading, this
+	// would lead to duplicate spans being "active" and forwarding the wrong one.
+	drop(_enter);
+	runtime.block_on(receiver).unwrap();
+	runtime.block_on(task_manager.clean_shutdown());
+	drop(runtime);
+
+	let spans = spans_entered.lock();
+	// We entered the telemetry span and the "test" in the future, the nested future and
+	// the "test" span outside of the future. So, we should have recorded 3 spans.
+	assert_eq!(5, spans.len());
+
+	assert_eq!(spans[0], "test");
+	assert_eq!(spans[1], telemetry_span.span().metadata().unwrap().name());
+	assert_eq!(spans[2], "test");
+	assert_eq!(spans[3], telemetry_span.span().metadata().unwrap().name());
+	assert_eq!(spans[4], "test");
+}