Runtime interface to add support for tracing from wasm (#6381)

* Add span recording to tracing implementation

* Add tracing proxy

* switch to rustc_hash::FxHashMap

* Replace lazy_static and hashmap with thread_local and vec.

* fix marking valid span as invalid while removing invalid spans

* refactor, add wasm_tracing module in `support`

* update registered spans

* tidy up

* typos

* refactor

* update flag name to signal lost trace - `is_valid_trace`

* update flag name to signal lost trace - `is_valid_trace`

* update docs

* update docs

* Use tracing Field recording to store the actual `name` and `target`
from wasm traces.

* fix debug log in subscriber + small refactor

* add tests

* handle misuse in case trying to exit span not held

* Implement filter for wasm traces, simplify field recording for primitive types

* remove superfluous warning

* update docs

* Update primitives/tracing/src/proxy.rs

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

* Apply suggestions from code review

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* update docs, apply suggestions

* move Proxy from thread_local to `Extension`, rename macro

* fix test

* unify native & wasm span macro calls

* implement wasm tracing control facility in primitives and frame

* add cli flag `--wasm-tracing`

* fix

* switch to `Option<u4>` (possible performance degradation), switch
to static mut bool

* performance improvement using u64 vs Option<u64>

* performance improvement moving concat to client

* update docs

* Update client/cli/src/params/import_params.rs

Co-authored-by: Cecile Tonglet <cecile@parity.io>

* performance improvement

* Revert "performance improvement"

This reverts commit 55ff8817a86302cd93bb6197eb4ca5bc7f4fb524.

* small refactor

* formatting

* bump impl_version

* Update client/cli/src/config.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* update docs

* small fixes, remove pub static

* nit

* add integration tests and refactor Subscriber

* tests

* revert formatting

* try fix test that works locally but not in CI

* try fix test that works locally but not in CI

* debug test that works locally but not in CI

* fix test that works locally but not in CI

* remove pub visibility from bool in runtime

* make TracingSpanGuard #[cfg(not(feature = "std"))], update docs, comments

* make TracingProxy drop implementation conditional on !empty state

* add docs for TraceHandler

* remove blank line

* update expect message

* update tests

* rename cli option to tracing_enable_wasm

* rename cli option to tracing_enable_wasm

* fix

* ensure wasm-tracing features are wasm only

* bump impl_version

* bump impl_version

* add `"pallet-scheduler/std"` to `[features]` `std` in node/runtime

* refactor service to remove sp_tracing dependency

* refactor: line width, trait bounds

* improve LogTraceHandler output

* fix test

* improve tracing log output

* Apply suggestions from code review

* Apply suggestions from code review

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* swap wasm indication from trace name to a separate value

* Update client/tracing/src/lib.rs

* add docs

* remove runtime features

remove wasm_tracing option from CLI

remove wasm_tracing flag from ProfilingSubscriber

Co-authored-by: Matt Rutherford <mattrutherford@users.noreply.github.com>
Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>
Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>
Co-authored-by: Cecile Tonglet <cecile@parity.io>

substrate/client/tracing/src/lib.rs

@@ -24,7 +24,7 @@
 //!
 //! Currently we provide `Log` (default), `Telemetry` variants for `Receiver`
 
-use std::collections::HashMap;
+use rustc_hash::FxHashMap;
 use std::fmt;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::{Duration, Instant};
@@ -38,10 +38,14 @@ use tracing_core::{
 	Level,
 	metadata::Metadata,
 	span::{Attributes, Id, Record},
-	subscriber::Subscriber
+	subscriber::Subscriber,
 };
 
 use sc_telemetry::{telemetry, SUBSTRATE_INFO};
+use sp_tracing::proxy::{WASM_NAME_KEY, WASM_TARGET_KEY, WASM_TRACE_IDENTIFIER};
+
+const ZERO_DURATION: Duration = Duration::from_nanos(0);
+const PROXY_TARGET: &'static str = "sp_tracing::proxy";
 
 /// Used to configure how to receive the metrics
 #[derive(Debug, Clone)]
@@ -58,36 +62,55 @@ impl Default for TracingReceiver {
 	}
 }
 
-#[derive(Debug)]
-struct SpanDatum {
-	id: u64,
-	name: &'static str,
-	target: &'static str,
-	level: Level,
-	line: u32,
-	start_time: Instant,
-	overall_time: Duration,
-	values: Visitor,
+/// A handler for tracing `SpanDatum`
+pub trait TraceHandler: Send + Sync {
+	/// Process a `SpanDatum`
+	fn process_span(&self, span: SpanDatum);
 }
 
+/// Represents a single instance of a tracing span
+#[derive(Debug)]
+pub struct SpanDatum {
+	pub id: u64,
+	pub name: String,
+	pub target: String,
+	pub level: Level,
+	pub line: u32,
+	pub start_time: Instant,
+	pub overall_time: Duration,
+	pub values: Visitor,
+}
+
+/// Holds associated values for a tracing span
 #[derive(Clone, Debug)]
-struct Visitor(Vec<(String, String)>);
+pub struct Visitor(FxHashMap<String, String>);
+
+impl Visitor {
+	/// Consume the Visitor, returning the inner FxHashMap
+	pub fn into_inner(self) -> FxHashMap<String, String> {
+		self.0
+	}
+}
 
 impl Visit for Visitor {
 	fn record_i64(&mut self, field: &Field, value: i64) {
-		self.record_debug(field, &value)
+		self.0.insert(field.name().to_string(), value.to_string());
 	}
 
 	fn record_u64(&mut self, field: &Field, value: u64) {
-		self.record_debug(field, &value)
+		self.0.insert(field.name().to_string(), value.to_string());
 	}
 
 	fn record_bool(&mut self, field: &Field, value: bool) {
-		self.record_debug(field, &value)
+		self.0.insert(field.name().to_string(), value.to_string());
+	}
+
+	fn record_str(&mut self, field: &Field, value: &str) {
+		self.0.insert(field.name().to_string(), value.to_owned());
 	}
 
 	fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
-		self.0.push((field.name().to_string(), format!("{:?}",value)));
+		self.0.insert(field.name().to_string(), format!("{:?}", value));
 	}
 }
 
@@ -105,7 +128,7 @@ impl Serialize for Visitor {
 
 impl fmt::Display for Visitor {
 	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-		let values = self.0.iter().map(|(k,v)| format!("{}={}",k,v)).collect::<Vec<String>>().join(", ");
+		let values = self.0.iter().map(|(k, v)| format!("{}={}", k, v)).collect::<Vec<String>>().join(", ");
 		write!(f, "{}", values)
 	}
 }
@@ -135,23 +158,50 @@ impl Value for Visitor {
 pub struct ProfilingSubscriber {
 	next_id: AtomicU64,
 	targets: Vec<(String, Level)>,
-	receiver: TracingReceiver,
-	span_data: Mutex<HashMap<u64, SpanDatum>>,
+	trace_handler: Box<dyn TraceHandler>,
+	span_data: Mutex<FxHashMap<u64, SpanDatum>>,
 }
 
 impl ProfilingSubscriber {
-	/// Takes a `Receiver` and a comma separated list of targets,
-	/// either with a level: "pallet=trace"
-	/// or without: "pallet".
-	pub fn new(receiver: TracingReceiver, targets: &str) -> Self {
+	/// Takes a `TracingReceiver` and a comma separated list of targets,
+	/// either with a level: "pallet=trace,frame=debug"
+	/// or without: "pallet,frame" in which case the level defaults to `trace`.
+	/// wasm_tracing indicates whether to enable wasm traces
+	pub fn new(receiver: TracingReceiver, targets: &str) -> ProfilingSubscriber {
+		match receiver {
+			TracingReceiver::Log => Self::new_with_handler(Box::new(LogTraceHandler), targets),
+			TracingReceiver::Telemetry => Self::new_with_handler(
+				Box::new(TelemetryTraceHandler),
+				targets,
+			),
+		}
+	}
+
+	/// Allows use of a custom TraceHandler to create a new instance of ProfilingSubscriber.
+	/// Takes a comma separated list of targets,
+	/// either with a level, eg: "pallet=trace"
+	/// or without: "pallet" in which case the level defaults to `trace`.
+	/// wasm_tracing indicates whether to enable wasm traces
+	pub fn new_with_handler(trace_handler: Box<dyn TraceHandler>, targets: &str)
+		-> ProfilingSubscriber
+	{
 		let targets: Vec<_> = targets.split(',').map(|s| parse_target(s)).collect();
 		ProfilingSubscriber {
 			next_id: AtomicU64::new(1),
 			targets,
-			receiver,
-			span_data: Mutex::new(HashMap::new()),
+			trace_handler,
+			span_data: Mutex::new(FxHashMap::default()),
 		}
 	}
+
+	fn check_target(&self, target: &str, level: &Level) -> bool {
+		for t in &self.targets {
+			if target.starts_with(t.0.as_str()) && level <= &t.1 {
+				return true;
+			}
+		}
+		false
+	}
 }
 
 // Default to TRACE if no level given or unable to parse Level
@@ -173,36 +223,45 @@ fn parse_target(s: &str) -> (String, Level) {
 
 impl Subscriber for ProfilingSubscriber {
 	fn enabled(&self, metadata: &Metadata<'_>) -> bool {
-		for t in &self.targets {
-			if metadata.target().starts_with(t.0.as_str()) && metadata.level() <= &t.1 {
-				log::debug!("Enabled target: {}, level: {}", metadata.target(), metadata.level());
-				return true;
-			} else {
-				log::debug!("Disabled target: {}, level: {}", metadata.target(), metadata.level());
-			}
+		if metadata.target() == PROXY_TARGET || self.check_target(metadata.target(), metadata.level()) {
+			log::debug!(target: "tracing", "Enabled target: {}, level: {}", metadata.target(), metadata.level());
+			true
+		} else {
+			log::debug!(target: "tracing", "Disabled target: {}, level: {}", metadata.target(), metadata.level());
+			false
 		}
-		false
 	}
 
 	fn new_span(&self, attrs: &Attributes<'_>) -> Id {
 		let id = self.next_id.fetch_add(1, Ordering::Relaxed);
-		let mut values = Visitor(Vec::new());
+		let mut values = Visitor(FxHashMap::default());
 		attrs.record(&mut values);
+		// If this is a wasm trace, check if target/level is enabled
+		if let Some(wasm_target) = values.0.get(WASM_TARGET_KEY) {
+			if !self.check_target(wasm_target, attrs.metadata().level()) {
+				return Id::from_u64(id);
+			}
+		}
 		let span_datum = SpanDatum {
 			id,
-			name: attrs.metadata().name(),
-			target: attrs.metadata().target(),
+			name: attrs.metadata().name().to_owned(),
+			target: attrs.metadata().target().to_owned(),
 			level: attrs.metadata().level().clone(),
 			line: attrs.metadata().line().unwrap_or(0),
 			start_time: Instant::now(),
-			overall_time: Duration::from_nanos(0),
+			overall_time: ZERO_DURATION,
 			values,
 		};
 		self.span_data.lock().insert(id, span_datum);
 		Id::from_u64(id)
 	}
 
-	fn record(&self, _span: &Id, _values: &Record<'_>) {}
+	fn record(&self, span: &Id, values: &Record<'_>) {
+		let mut span_data = self.span_data.lock();
+		if let Some(s) = span_data.get_mut(&span.into_u64()) {
+			values.record(&mut s.values);
+		}
+	}
 
 	fn record_follows_from(&self, _span: &Id, _follows: &Id) {}
 
@@ -213,65 +272,89 @@ impl Subscriber for ProfilingSubscriber {
 		let start_time = Instant::now();
 		if let Some(mut s) = span_data.get_mut(&span.into_u64()) {
 			s.start_time = start_time;
-		} else {
-			log::warn!("Tried to enter span {:?} that has already been closed!", span);
 		}
 	}
 
 	fn exit(&self, span: &Id) {
-		let mut span_data = self.span_data.lock();
 		let end_time = Instant::now();
+		let mut span_data = self.span_data.lock();
 		if let Some(mut s) = span_data.get_mut(&span.into_u64()) {
 			s.overall_time = end_time - s.start_time + s.overall_time;
 		}
 	}
 
 	fn try_close(&self, span: Id) -> bool {
-		let mut span_data = self.span_data.lock();
-		if let Some(data) = span_data.remove(&span.into_u64()) {
-			self.send_span(data);
+		let span_datum = {
+			let mut span_data = self.span_data.lock();
+			span_data.remove(&span.into_u64())
+		};
+		if let Some(mut span_datum) = span_datum {
+			if span_datum.name == WASM_TRACE_IDENTIFIER {
+				span_datum.values.0.insert("wasm".to_owned(), "true".to_owned());
+				if let Some(n) = span_datum.values.0.remove(WASM_NAME_KEY) {
+					span_datum.name = n;
+				}
+				if let Some(t) = span_datum.values.0.remove(WASM_TARGET_KEY) {
+					span_datum.target = t;
+				}
+			}
+			if self.check_target(&span_datum.target, &span_datum.level) {
+				self.trace_handler.process_span(span_datum);
+			}
 		};
 		true
 	}
 }
 
-impl ProfilingSubscriber {
-	fn send_span(&self, span_datum: SpanDatum) {
-		match self.receiver {
-			TracingReceiver::Log => print_log(span_datum),
-			TracingReceiver::Telemetry => send_telemetry(span_datum),
+/// TraceHandler for sending span data to the logger
+pub struct LogTraceHandler;
+
+fn log_level(level: Level) -> log::Level {
+	match level {
+		Level::TRACE => log::Level::Trace,
+		Level::DEBUG => log::Level::Debug,
+		Level::INFO => log::Level::Info,
+		Level::WARN => log::Level::Warn,
+		Level::ERROR => log::Level::Error,
+	}
+}
+
+impl TraceHandler for LogTraceHandler {
+	fn process_span(&self, span_datum: SpanDatum) {
+		if span_datum.values.0.is_empty() {
+			log::log!(
+				log_level(span_datum.level), 
+				"{}: {}, time: {}",
+				span_datum.target,
+				span_datum.name,
+				span_datum.overall_time.as_nanos(),
+			);
+		} else {
+			log::log!(
+				log_level(span_datum.level),
+				"{}: {}, time: {}, {}",
+				span_datum.target,
+				span_datum.name,
+				span_datum.overall_time.as_nanos(),
+				span_datum.values,
+			);
 		}
 	}
 }
 
-fn print_log(span_datum: SpanDatum) {
-	if span_datum.values.0.is_empty() {
-		log::info!("TRACING: {} {}: {}, line: {}, time: {}",
-			span_datum.level,
-			span_datum.target,
-			span_datum.name,
-			span_datum.line,
-			span_datum.overall_time.as_nanos(),
-		);
-	} else {
-		log::info!("TRACING: {} {}: {}, line: {}, time: {}, {}",
-			span_datum.level,
-			span_datum.target,
-			span_datum.name,
-			span_datum.line,
-			span_datum.overall_time.as_nanos(),
-			span_datum.values
+/// TraceHandler for sending span data to telemetry,
+/// Please see telemetry documentation for details on how to specify endpoints and
+/// set the required telemetry level to activate tracing messages
+pub struct TelemetryTraceHandler;
+
+impl TraceHandler for TelemetryTraceHandler {
+	fn process_span(&self, span_datum: SpanDatum) {
+		telemetry!(SUBSTRATE_INFO; "tracing.profiling";
+			"name" => span_datum.name,
+			"target" => span_datum.target,
+			"line" => span_datum.line,
+			"time" => span_datum.overall_time.as_nanos(),
+			"values" => span_datum.values
 		);
 	}
 }
-
-fn send_telemetry(span_datum: SpanDatum) {
-	telemetry!(SUBSTRATE_INFO; "tracing.profiling";
-		"name" => span_datum.name,
-		"target" => span_datum.target,
-		"line" => span_datum.line,
-		"time" => span_datum.overall_time.as_nanos(),
-		"values" => span_datum.values
-	);
-}
-