Add logger configuration hook (#10440)

* Initial poc

Signed-off-by: Andrei Sandu <andrei-mihail@parity.io>

* Make config available to logger hook

Signed-off-by: Andrei Sandu <andrei-mihail@parity.io>

* fmt

Signed-off-by: Andrei Sandu <andrei-mihail@parity.io>

* Fix tests

Signed-off-by: Andrei Sandu <andrei-mihail@parity.io>

* fmt

Signed-off-by: Andrei Sandu <andrei-mihail@parity.io>

* Add metric prefix option in sc_cli::RunCmd

Signed-off-by: Andrei Sandu <andrei-mihail@parity.io>

* Remove print

Signed-off-by: Andrei Sandu <andrei-mihail@parity.io>

* Review fixes

Signed-off-by: Andrei Sandu <andrei-mihail@parity.io>

* fmt

Signed-off-by: Andrei Sandu <andrei-mihail@parity.io>

* fix docs

Signed-off-by: Andrei Sandu <andrei-mihail@parity.io>

client/cli/src/config.rs

@@ -581,10 +581,40 @@ pub trait CliConfiguration<DCV: DefaultConfigurationValues = ()>: Sized {
 	/// This method:
 	///
 	/// 1. Sets the panic handler
+	/// 2. Optionally customize logger/profiling
 	/// 2. Initializes the logger
 	/// 3. Raises the FD limit
-	fn init<C: SubstrateCli>(&self) -> Result<()> {
-		sp_panic_handler::set(&C::support_url(), &C::impl_version());
+	///
+	/// The `logger_hook` closure is executed before the logger is constructed
+	/// and initialized. It is useful for setting up a custom profiler.
+	///
+	/// Example:
+	/// ```
+	/// use sc_tracing::{SpanDatum, TraceEvent};
+	/// struct TestProfiler;
+	///
+	/// impl sc_tracing::TraceHandler for TestProfiler {
+	///  	fn handle_span(&self, sd: &SpanDatum) {}
+	/// 		fn handle_event(&self, _event: &TraceEvent) {}
+	/// };
+	///
+	/// fn logger_hook() -> impl FnOnce(&mut sc_cli::LoggerBuilder, &sc_service::Configuration) -> () {
+	/// 	|logger_builder, config| {
+	/// 			logger_builder.with_custom_profiling(Box::new(TestProfiler{}));
+	/// 	}
+	/// }
+	/// ```
+	fn init<F>(
+		&self,
+		support_url: &String,
+		impl_version: &String,
+		logger_hook: F,
+		config: &Configuration,
+	) -> Result<()>
+	where
+		F: FnOnce(&mut LoggerBuilder, &Configuration),
+	{
+		sp_panic_handler::set(support_url, impl_version);
 
 		let mut logger = LoggerBuilder::new(self.log_filters()?);
 		logger
@@ -600,6 +630,9 @@ pub trait CliConfiguration<DCV: DefaultConfigurationValues = ()>: Sized {
 			logger.with_colors(false);
 		}
 
+		// Call hook for custom profiling setup.
+		logger_hook(&mut logger, &config);
+
 		logger.init()?;
 
 		if let Some(new_limit) = fdlimit::raise_fd_limit() {

client/cli/src/lib.rs

@@ -224,10 +224,46 @@ pub trait SubstrateCli: Sized {
 	/// Create a runner for the command provided in argument. This will create a Configuration and
 	/// a tokio runtime
 	fn create_runner<T: CliConfiguration>(&self, command: &T) -> error::Result<Runner<Self>> {
-		command.init::<Self>()?;
-		Runner::new(self, command)
+		let tokio_runtime = build_runtime()?;
+		let config = command.create_configuration(self, tokio_runtime.handle().clone())?;
+
+		command.init(&Self::support_url(), &Self::impl_version(), |_, _| {}, &config)?;
+		Runner::new(config, tokio_runtime)
 	}
 
+	/// Create a runner for the command provided in argument. The `logger_hook` can be used to setup
+	/// a custom profiler or update the logger configuration before it is initialized.
+	///
+	/// Example:
+	/// ```
+	/// use sc_tracing::{SpanDatum, TraceEvent};
+	/// struct TestProfiler;
+	///
+	/// impl sc_tracing::TraceHandler for TestProfiler {
+	///  	fn handle_span(&self, sd: &SpanDatum) {}
+	/// 		fn handle_event(&self, _event: &TraceEvent) {}
+	/// };
+	///
+	/// fn logger_hook() -> impl FnOnce(&mut sc_cli::LoggerBuilder, &sc_service::Configuration) -> () {
+	/// 	|logger_builder, config| {
+	/// 			logger_builder.with_custom_profiling(Box::new(TestProfiler{}));
+	/// 	}
+	/// }
+	/// ```
+	fn create_runner_with_logger_hook<T: CliConfiguration, F>(
+		&self,
+		command: &T,
+		logger_hook: F,
+	) -> error::Result<Runner<Self>>
+	where
+		F: FnOnce(&mut LoggerBuilder, &Configuration),
+	{
+		let tokio_runtime = build_runtime()?;
+		let config = command.create_configuration(self, tokio_runtime.handle().clone())?;
+
+		command.init(&Self::support_url(), &Self::impl_version(), logger_hook, &config)?;
+		Runner::new(config, tokio_runtime)
+	}
 	/// Native runtime version.
 	fn native_runtime_version(chain_spec: &Box<dyn ChainSpec>) -> &'static RuntimeVersion;
 }

client/cli/src/runner.rs

@@ -16,7 +16,7 @@
 // You should have received a copy of the GNU General Public License
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
-use crate::{error::Error as CliError, CliConfiguration, Result, SubstrateCli};
+use crate::{error::Error as CliError, Result, SubstrateCli};
 use chrono::prelude::*;
 use futures::{future, future::FutureExt, pin_mut, select, Future};
 use log::info;
@@ -112,15 +112,8 @@ pub struct Runner<C: SubstrateCli> {
 
 impl<C: SubstrateCli> Runner<C> {
 	/// Create a new runtime with the command provided in argument
-	pub fn new<T: CliConfiguration>(cli: &C, command: &T) -> Result<Runner<C>> {
-		let tokio_runtime = build_runtime()?;
-		let runtime_handle = tokio_runtime.handle().clone();
-
-		Ok(Runner {
-			config: command.create_configuration(cli, runtime_handle)?,
-			tokio_runtime,
-			phantom: PhantomData,
-		})
+	pub fn new(config: Configuration, tokio_runtime: tokio::runtime::Runtime) -> Result<Runner<C>> {
+		Ok(Runner { config, tokio_runtime, phantom: PhantomData })
 	}
 
 	/// Log information about the node itself.

client/executor/src/integration_tests/mod.rs

@@ -622,11 +622,11 @@ fn wasm_tracing_should_work(wasm_method: WasmExecutionMethod) {
 	struct TestTraceHandler(Arc<Mutex<Vec<SpanDatum>>>);
 
 	impl sc_tracing::TraceHandler for TestTraceHandler {
-		fn handle_span(&self, sd: SpanDatum) {
-			self.0.lock().unwrap().push(sd);
+		fn handle_span(&self, sd: &SpanDatum) {
+			self.0.lock().unwrap().push(sd.clone());
 		}
 
-		fn handle_event(&self, _event: TraceEvent) {}
+		fn handle_event(&self, _event: &TraceEvent) {}
 	}
 
 	let traces = Arc::new(Mutex::new(Vec::new()));

client/service/src/config.rs

@@ -188,8 +188,6 @@ pub struct PrometheusConfig {
 
 impl PrometheusConfig {
 	/// Create a new config using the default registry.
-	///
-	/// The default registry prefixes metrics with `substrate`.
 	pub fn new_with_default_registry(port: SocketAddr, chain_id: String) -> Self {
 		let param = iter::once((String::from("chain"), chain_id)).collect();
 		Self {

client/tracing/src/lib.rs

@@ -58,7 +58,7 @@ const ZERO_DURATION: Duration = Duration::from_nanos(0);
 /// Responsible for assigning ids to new spans, which are not re-used.
 pub struct ProfilingLayer {
 	targets: Vec<(String, Level)>,
-	trace_handler: Box<dyn TraceHandler>,
+	trace_handlers: Vec<Box<dyn TraceHandler>>,
 }
 
 /// Used to configure how to receive the metrics
@@ -76,14 +76,14 @@ impl Default for TracingReceiver {
 
 /// A handler for tracing `SpanDatum`
 pub trait TraceHandler: Send + Sync {
-	/// Process a `SpanDatum`
-	fn handle_span(&self, span: SpanDatum);
-	/// Process a `TraceEvent`
-	fn handle_event(&self, event: TraceEvent);
+	/// Process a `SpanDatum`.
+	fn handle_span(&self, span: &SpanDatum);
+	/// Process a `TraceEvent`.
+	fn handle_event(&self, event: &TraceEvent);
 }
 
 /// Represents a tracing event, complete with values
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct TraceEvent {
 	/// Name of the event.
 	pub name: String,
@@ -98,7 +98,7 @@ pub struct TraceEvent {
 }
 
 /// Represents a single instance of a tracing span
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct SpanDatum {
 	/// id for this span
 	pub id: Id,
@@ -213,6 +213,15 @@ impl fmt::Display for Values {
 	}
 }
 
+/// Trace handler event types.
+#[derive(Debug)]
+pub enum TraceHandlerEvents {
+	/// An event.
+	Event(TraceEvent),
+	/// A span.
+	Span(SpanDatum),
+}
+
 impl ProfilingLayer {
 	/// Takes a `TracingReceiver` and a comma separated list of targets,
 	/// either with a level: "pallet=trace,frame=debug"
@@ -231,7 +240,12 @@ impl ProfilingLayer {
 	/// wasm_tracing indicates whether to enable wasm traces
 	pub fn new_with_handler(trace_handler: Box<dyn TraceHandler>, targets: &str) -> Self {
 		let targets: Vec<_> = targets.split(',').map(|s| parse_target(s)).collect();
-		Self { targets, trace_handler }
+		Self { targets, trace_handlers: vec![trace_handler] }
+	}
+
+	/// Attach additional handlers to allow handling of custom events/spans.
+	pub fn add_handler(&mut self, trace_handler: Box<dyn TraceHandler>) {
+		self.trace_handlers.push(trace_handler);
 	}
 
 	fn check_target(&self, target: &str, level: &Level) -> bool {
@@ -242,6 +256,18 @@ impl ProfilingLayer {
 		}
 		false
 	}
+
+	/// Sequentially dispatch a trace event to all handlers.
+	fn dispatch_event(&self, event: TraceHandlerEvents) {
+		match &event {
+			TraceHandlerEvents::Span(span_datum) => {
+				self.trace_handlers.iter().for_each(|handler| handler.handle_span(span_datum));
+			},
+			TraceHandlerEvents::Event(event) => {
+				self.trace_handlers.iter().for_each(|handler| handler.handle_event(event));
+			},
+		}
+	}
 }
 
 // Default to TRACE if no level given or unable to parse Level
@@ -320,7 +346,7 @@ where
 			values,
 			parent_id,
 		};
-		self.trace_handler.handle_event(trace_event);
+		self.dispatch_event(TraceHandlerEvents::Event(trace_event));
 	}
 
 	fn on_enter(&self, span: &Id, ctx: Context<S>) {
@@ -348,10 +374,10 @@ where
 						span_datum.target = t;
 					}
 					if self.check_target(&span_datum.target, &span_datum.level) {
-						self.trace_handler.handle_span(span_datum);
+						self.dispatch_event(TraceHandlerEvents::Span(span_datum));
 					}
 				} else {
-					self.trace_handler.handle_span(span_datum);
+					self.dispatch_event(TraceHandlerEvents::Span(span_datum));
 				}
 			}
 		}
@@ -374,7 +400,7 @@ fn log_level(level: Level) -> log::Level {
 }
 
 impl TraceHandler for LogTraceHandler {
-	fn handle_span(&self, span_datum: SpanDatum) {
+	fn handle_span(&self, span_datum: &SpanDatum) {
 		if span_datum.values.is_empty() {
 			log::log!(
 				log_level(span_datum.level),
@@ -383,7 +409,7 @@ impl TraceHandler for LogTraceHandler {
 				span_datum.name,
 				span_datum.overall_time.as_nanos(),
 				span_datum.id.into_u64(),
-				span_datum.parent_id.map(|s| s.into_u64()),
+				span_datum.parent_id.as_ref().map(|s| s.into_u64()),
 			);
 		} else {
 			log::log!(
@@ -393,18 +419,18 @@ impl TraceHandler for LogTraceHandler {
 				span_datum.name,
 				span_datum.overall_time.as_nanos(),
 				span_datum.id.into_u64(),
-				span_datum.parent_id.map(|s| s.into_u64()),
+				span_datum.parent_id.as_ref().map(|s| s.into_u64()),
 				span_datum.values,
 			);
 		}
 	}
 
-	fn handle_event(&self, event: TraceEvent) {
+	fn handle_event(&self, event: &TraceEvent) {
 		log::log!(
 			log_level(event.level),
 			"{}, parent_id: {:?}, {}",
 			event.target,
-			event.parent_id.map(|s| s.into_u64()),
+			event.parent_id.as_ref().map(|s| s.into_u64()),
 			event.values,
 		);
 	}
@@ -447,12 +473,12 @@ mod tests {
 	}
 
 	impl TraceHandler for TestTraceHandler {
-		fn handle_span(&self, sd: SpanDatum) {
-			self.spans.lock().push(sd);
+		fn handle_span(&self, sd: &SpanDatum) {
+			self.spans.lock().push(sd.clone());
 		}
 
-		fn handle_event(&self, event: TraceEvent) {
-			self.events.lock().push(event);
+		fn handle_event(&self, event: &TraceEvent) {
+			self.events.lock().push(event.clone());
 		}
 	}
 

client/tracing/src/logging/mod.rs

@@ -193,6 +193,7 @@ where
 pub struct LoggerBuilder {
 	directives: String,
 	profiling: Option<(crate::TracingReceiver, String)>,
+	custom_profiler: Option<Box<dyn crate::TraceHandler>>,
 	log_reloading: bool,
 	force_colors: Option<bool>,
 	detailed_output: bool,
@@ -204,6 +205,7 @@ impl LoggerBuilder {
 		Self {
 			directives: directives.into(),
 			profiling: None,
+			custom_profiler: None,
 			log_reloading: false,
 			force_colors: None,
 			detailed_output: false,
@@ -220,6 +222,15 @@ impl LoggerBuilder {
 		self
 	}
 
+	/// Add a custom profiler.
+	pub fn with_custom_profiling(
+		&mut self,
+		custom_profiler: Box<dyn crate::TraceHandler>,
+	) -> &mut Self {
+		self.custom_profiler = Some(custom_profiler);
+		self
+	}
+
 	/// Wether or not to disable log reloading.
 	pub fn with_log_reloading(&mut self, enabled: bool) -> &mut Self {
 		self.log_reloading = enabled;
@@ -256,7 +267,12 @@ impl LoggerBuilder {
 					self.detailed_output,
 					|builder| enable_log_reloading!(builder),
 				)?;
-				let profiling = crate::ProfilingLayer::new(tracing_receiver, &profiling_targets);
+				let mut profiling =
+					crate::ProfilingLayer::new(tracing_receiver, &profiling_targets);
+
+				self.custom_profiler
+					.into_iter()
+					.for_each(|profiler| profiling.add_handler(profiler));
 
 				tracing::subscriber::set_global_default(subscriber.with(profiling))?;
 
@@ -269,7 +285,12 @@ impl LoggerBuilder {
 					self.detailed_output,
 					|builder| builder,
 				)?;
-				let profiling = crate::ProfilingLayer::new(tracing_receiver, &profiling_targets);
+				let mut profiling =
+					crate::ProfilingLayer::new(tracing_receiver, &profiling_targets);
+
+				self.custom_profiler
+					.into_iter()
+					.for_each(|profiler| profiling.add_handler(profiler));
 
 				tracing::subscriber::set_global_default(subscriber.with(profiling))?;