Add log spans using thread local storage in std

Author: joostjager

lightning-macros/src/lib.rs

@@ -20,12 +20,13 @@
 
 extern crate alloc;
 
+use alloc::string::String;
 use alloc::string::ToString;
 use proc_macro::{Delimiter, Group, TokenStream, TokenTree};
 use proc_macro2::TokenStream as TokenStream2;
 use quote::quote;
 use syn::spanned::Spanned;
-use syn::{parse, ImplItemFn, Token};
+use syn::{parse, parse2, Expr, ExprLit, ImplItemFn, Lit, Token};
 use syn::{parse_macro_input, Item};
 
 fn add_async_method(mut parsed: ImplItemFn) -> TokenStream {
@@ -400,3 +401,44 @@ pub fn xtest_inventory(_input: TokenStream) -> TokenStream {
 
 	TokenStream::from(expanded)
 }
+
+/// Adds a logging scope at the top of a method.
+#[proc_macro_attribute]
+pub fn log_scope(attrs: TokenStream, meth: TokenStream) -> TokenStream {
+	let attrs: TokenStream2 = parse_macro_input!(attrs as TokenStream2);
+	let mut name_attr: Option<String> = None;
+	if !attrs.is_empty() {
+		// Expect something like `name = "foo"`
+		let expr: Expr = parse2(attrs.clone()).expect("invalid attribute syntax");
+
+		if let Expr::Assign(assign) = expr {
+			// Check left-hand side is `name`.
+			if let Expr::Path(path) = *assign.left {
+				if path.path.is_ident("name") {
+					if let Expr::Lit(ExprLit { lit: Lit::Str(s), .. }) = *assign.right {
+						name_attr = Some(s.value());
+					}
+				}
+			}
+		}
+	}
+
+	let mut meth = if let Ok(parsed) = parse::<syn::ItemFn>(meth) {
+		parsed
+	} else {
+		return (quote! {
+			compile_error!("log_scope can only be set on methods")
+		})
+		.into();
+	};
+
+	// Use the attribute name if present, otherwise fall back to the function name
+	let name = name_attr.unwrap_or_else(|| meth.sig.ident.to_string());
+
+	let stmt = quote! {
+		let _logging_context = crate::util::logger::LoggerScope::new(#name);
+	};
+
+	meth.block.stmts.insert(0, parse(stmt.into()).unwrap());
+	quote! { #meth }.into()
+}

lightning/src/util/logger.rs

@@ -15,6 +15,7 @@
 
 use bitcoin::secp256k1::PublicKey;
 
+use core::cell::RefCell;
 use core::cmp;
 use core::fmt;
 use core::fmt::Display;
@@ -130,6 +131,9 @@ pub struct Record<$($args)?> {
 	/// Note that this is only filled in for logs pertaining to a specific payment, and will be
 	/// `None` for logs which are not directly related to a payment.
 	pub payment_hash: Option<PaymentHash>,
+	/// The names of the surrounding spans, if any.
+	#[cfg(feature = "std")]
+	pub spans: Vec<&'static str>,
 }
 
 impl<$($args)?> Record<$($args)?> {
@@ -138,7 +142,10 @@ impl<$($args)?> Record<$($args)?> {
 	/// This is not exported to bindings users as fmt can't be used in C
 	#[inline]
 	pub fn new<$($nonstruct_args)?>(
-		level: Level, peer_id: Option<PublicKey>, channel_id: Option<ChannelId>,
+		level: Level,
+		#[cfg(feature = "std")]
+		spans: Vec<&'static str>,
+		peer_id: Option<PublicKey>, channel_id: Option<ChannelId>,
 		args: fmt::Arguments<'a>, module_path: &'static str, file: &'static str, line: u32,
 		payment_hash: Option<PaymentHash>
 	) -> Record<$($args)?> {
@@ -154,6 +161,8 @@ impl<$($args)?> Record<$($args)?> {
 			file,
 			line,
 			payment_hash,
+			#[cfg(feature = "std")]
+			spans,
 		}
 	}
 }
@@ -189,7 +198,20 @@ impl<$($args)?> Display for Record<$($args)?> {
 
 		#[cfg(test)]
 		{
-			write!(f, " {}", self.args)?;
+			write!(f, " ")?;
+			#[cfg(feature = "std")]
+			if !self.spans.is_empty() {
+				write!(f, "[")?;
+				for (i, span) in self.spans.iter().enumerate() {
+					if i > 0 {
+						write!(f, "->")?;
+					}
+					write!(f, "{}", span)?;
+				}
+				write!(f, "] ")?;
+			}
+
+			write!(f, "{}", self.args)?;
 
 			let mut open_bracket_written = false;
 			if let Some(peer_id) = self.peer_id {
@@ -384,14 +406,70 @@ impl<T: fmt::Display, I: core::iter::Iterator<Item = T> + Clone> fmt::Display fo
 	}
 }
 
+#[cfg(feature = "std")]
+thread_local! {
+	/// The thread-local stack of loggers.
+	pub static TLS_LOGGER: RefCell<Vec<&'static str>> = const { RefCell::new(Vec::new()) };
+}
+
+/// A scope which pushes a logger on a thread-local stack for the duration of the scope.
+pub struct LoggerScope<'a> {
+	_marker: core::marker::PhantomData<&'a ()>,
+}
+
+impl<'a> LoggerScope<'a> {
+	/// Pushes a logger onto the thread-local logger stack.
+	pub fn new(span: &'static str) -> Self {
+		#[cfg(feature = "std")]
+		TLS_LOGGER.with(|cell| {
+			let mut stack = cell.borrow_mut();
+			stack.push(span);
+		});
+
+		LoggerScope { _marker: core::marker::PhantomData }
+	}
+}
+
+impl<'a> Drop for LoggerScope<'a> {
+	fn drop(&mut self) {
+		#[cfg(feature = "std")]
+		TLS_LOGGER.with(|cell| {
+			let mut stack = cell.borrow_mut();
+			stack.pop();
+		});
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use crate::ln::types::ChannelId;
 	use crate::sync::Arc;
 	use crate::types::payment::PaymentHash;
-	use crate::util::logger::{Level, Logger, WithContext};
+	use crate::util::logger::{Level, Logger, LoggerScope, WithContext};
 	use crate::util::test_utils::TestLogger;
 	use bitcoin::secp256k1::{PublicKey, Secp256k1, SecretKey};
+	use lightning_macros::log_scope;
+
+	#[test]
+	fn logger_scope() {
+		let logger = TestLogger::new();
+		let _scope = LoggerScope::new("test_logger_scope");
+		log_info!(logger, "Info");
+	}
+
+	#[test]
+	#[log_scope(name = "test_logger_scope_proc_macro")]
+	fn logger_scope_proc_macro() {
+		let logger = TestLogger::new();
+		log_info!(logger, "Info");
+	}
+
+	#[test]
+	#[log_scope]
+	fn logger_scope_proc_macro_no_name() {
+		let logger = TestLogger::new();
+		log_info!(logger, "Info");
+	}
 
 	#[test]
 	fn test_level_show() {

lightning/src/util/macro_logger.rs

@@ -174,9 +174,17 @@ macro_rules! log_spendable {
 #[doc(hidden)]
 #[macro_export]
 macro_rules! log_given_level {
-	($logger: expr, $lvl:expr, $($arg:tt)+) => (
-		$logger.log($crate::util::logger::Record::new($lvl, None, None, format_args!($($arg)+), module_path!(), file!(), line!(), None))
-	);
+	($logger: expr, $lvl:expr, $($arg:tt)+) => {{
+		#[cfg(feature = "std")]
+		let spans = $crate::util::logger::TLS_LOGGER.with(|cell| {
+			cell.borrow().iter().map(|span| *span).collect()
+		});
+		$logger.log($crate::util::logger::Record::new($lvl,
+			#[cfg(feature = "std")]
+			spans,
+			None, None, format_args!($($arg)+), module_path!(), file!(), line!(), None
+		));
+	}};
 }
 
 /// Log at the `ERROR` level.

lightning/src/util/test_utils.rs

@@ -55,12 +55,15 @@ use crate::util::config::UserConfig;
 use crate::util::dyn_signer::{
 	DynKeysInterface, DynKeysInterfaceTrait, DynPhantomKeysInterface, DynSigner,
 };
+use crate::util::logger::LoggerScope;
 use crate::util::logger::{Logger, Record};
 #[cfg(feature = "std")]
 use crate::util::mut_global::MutGlobal;
 use crate::util::persist::{KVStore, KVStoreSync, MonitorName};
 use crate::util::ser::{Readable, ReadableArgs, Writeable, Writer};
 use crate::util::test_channel_signer::{EnforcementState, TestChannelSigner};
+#[cfg(feature = "std")]
+use std::cell::RefCell;
 
 use bitcoin::amount::Amount;
 use bitcoin::block::Block;
@@ -1759,6 +1762,26 @@ impl Logger for TestLogger {
 	}
 }
 
+#[cfg(feature = "std")]
+thread_local! {
+	pub(crate) static THREAD_LOG_SCOPE: RefCell<Option<LoggerScope<'static>>> = RefCell::new(None);
+}
+
+/// Sets up a logging scope for the current thread with the given span name. This is useful to split up a long (test) function
+/// into multiple scopes for easier log analysis.
+#[macro_export]
+macro_rules! test_scope {
+	($span:expr) => {
+		#[cfg(feature = "std")]
+		crate::util::test_utils::THREAD_LOG_SCOPE.with(|cell| {
+			// Drop old scope if it exists.
+			let _ = cell.borrow_mut().take();
+			// Create new scope.
+			*cell.borrow_mut() = Some(crate::util::logger::LoggerScope::new($span));
+		});
+	};
+}
+
 pub struct TestNodeSigner {
 	node_secret: SecretKey,
 }