node spans

Author: joostjager

lightning-macros/src/lib.rs

@@ -417,7 +417,8 @@ pub fn auto_span_methods(_attr: TokenStream, item: TokenStream) -> TokenStream {
 				let block = &method.block;
 				let method_name = method.sig.ident.to_string();
 				method.block = syn::parse_quote!({
-					let _entered_span = tracing::span!(tracing::Level::INFO, #method_name, node_id = self.node_id).entered();
+					let _node_span = enter_instance_span(self.node_id.clone()).entered();
+					let _method_span = tracing::span!(tracing::Level::INFO, #method_name).entered();
 					#block
 				});
 			}

lightning/src/ln/channelmanager.rs

@@ -31,6 +31,7 @@ use bitcoin::hashes::{Hash, HashEngine, HmacEngine};
 use bitcoin::secp256k1::Secp256k1;
 use bitcoin::secp256k1::{PublicKey, SecretKey};
 use bitcoin::{secp256k1, Sequence, SignedAmount};
+use tracing::Span;
 
 use crate::blinded_path::message::{
 	AsyncPaymentsContext, BlindedMessagePath, MessageForwardNode, OffersContext,
@@ -1889,6 +1890,36 @@ where
 	}
 }
 
+thread_local! {
+	static CURRENT_INSTANCE: RefCell<Option<String>> = RefCell::new(None);
+	static CURRENT_SPAN: RefCell<Option<Span>> = RefCell::new(None);
+}
+
+fn enter_instance_span(name: String) -> tracing::span::Span {
+	let mut span_to_enter: Option<Span> = None;
+
+	CURRENT_INSTANCE.with(|current_name| {
+		let mut current_name = current_name.borrow_mut();
+
+		if current_name.as_deref() != Some(&name) {
+			// Start a new span and store it
+			let span = tracing::info_span!("node", name);
+			CURRENT_SPAN.with(|s| *s.borrow_mut() = Some(span.clone())); // clone is cheap
+			*current_name = Some(name);
+			span_to_enter = Some(span);
+		} else {
+			// Reuse existing span
+			CURRENT_SPAN.with(|s| {
+				if let Some(span) = s.borrow().as_ref() {
+					span_to_enter = Some(span.clone());
+				}
+			});
+		}
+	});
+
+	span_to_enter.unwrap()
+}
+
 /// A lightning node's channel state machine and payment management logic, which facilitates
 /// sending, forwarding, and receiving payments through lightning channels.
 ///