Consolidate subsystem spans so they are all children of the leaf-activated root span (#6458)

* Pass the PerLeafSpan as mutable reference to handle_new_head function

* cargo +nightly fmt --all

* Add mock span for test

* cargo +nightly fmt --all

* add new-blocks-hashes to span

* ref span in match statement, set span to disabled if not passed

* remove second match clause, make handle_new_head_span mutable

* cargo +nightly fmt --all

* improve tag on error and warning

* add imported blocks and info span

* cargo +nightly fmt --all

* Improve error for imported_blocks_and_info trace

* format tags on get_header_span

* add lost-to-finality tag

* add missing bracket

* - Add bitfield child span
- Add block db insertion span

* - fix update-bitfield span tag

* - Fix type conversion to u64
- Add missing argument

* - Cargo fmt

* - Test add_follows_from

* - Revert as  relationship between spans not working correctly

* - use drop to test if parent-child relationship can be re-established

* - remove bitfield span, check if parent-child relationship can be reestablished

* - Remove dangling bitfield span which is not used, to see if parent-child relationship can be re-established

* Another dangling bitfield span

* cargo fmt

* - add imported blocks and info span
- add candidate span per candidate

* add tags before moving block_header to push scope

* - Add db-insertion span

* cargo fmt

* fix types

* * Pass mutable reference to span in handle_new_head
* Change get-header-span tags in handle_new_head
* Create cache-session-info span in handle_new_head
* Create optional argument in determine_new_blocks
* Pass mutable reference to handle_new_head_span in determine_new_blocks in handle_new_head function
* Add candidate-hash, candidate-number, lost-to-finality tags to candidate_span in handle_new_head function
* Manually drop db_insertion_span and remove superfluous tags  to it, only keeping approved-bitfields tag
* Add ApprovalVoting stage in jaeger

* * Pass mutable reference to jaeger::Span in stead of PerLeafSpan
* Add block-import span

* *Pass optional_span (optional argument) to determine_new_blocks util function

* * Add num-candidates int tag to block_import_span

* * Add head tag to cache_session_span

* * Create PerLeafSpan in handle_from_overseer (this is required to establish parent-child relationship between approval-voting span, and leaf-activated root span)

* * Add candidate-import-span as child of block-import-span
* Add candidate-hash and num-approval tags to candidate-import-span

* * Fix num-candidate tag to bitvec-len tag in candidate-import-span

* *Fix imported_blocKs_and_info span to create new-block-span as not dealing with candidates

* Consider the future::select! block

* Use HashMap<Hash, jaeger::PerLeafSpan>

* Remove Stage 9

* Add missing spans

* cargo +nightly fmt --all

* Remove optional span argument for determine_new_blocks

* * Remove no-longer needed default PerLeafSpan implementation
* Remove no-longer necessary mock span given re-factoring of handle_new_head() no longer neeing mutable span
* Split validation-result and request-data (availability and validation code) spans into two by dropping request_validation_data_spans
* Remove drop statements for cache_session_info_span
*

* Remove unnecessary span

* Remove another excessively spammy span

* Add missing spans from State in import tests

* Use functional approach to get spans

* - Add functional approach for the approval-voting span
- Add doc on block_numbers given labelling ambiguity
- Add span pruning logic
- Use .add_para_id on validation_result_span

* Replace for hash_set in hash_set_iter with map closure

* cargo +nightly fmt --all

* Change from unconsumed `map` to `.for_each`

* cargo +nightly fmt --all

* Refactor add_para_id to validation_result_span

* cargo +nightly fmt --all

* Remove duplicate tag

* Add missing tag to handle-approved-ancestor span

* Refactor span pruning to only invoke retain once

* Typo in span name

* - Replace unwrap_or with unwrap_or_else due to lazy evaluation of trace-identifier in polkadot_node_jaeger
- Remove some redundant spans

* Add approval-distribution spans

* - Add unwrap_or_else on note-approved-in-chain-selection
- Use child_with_trace_id to add traceID string tag on span (note this does not change the traceID, but just adds a tag)

* cargo +nightly fmt --all

* - Add traceID tags were necessary in approval-voting and availability-distribution
- Always use block-hash tag in stead of relay-parent tag in approval-distribution

* Remove schedule-wakeup span as it will duplicate spans on existing wakeups (which should be a no-op)

* Remove a couple of warnings related to mutability

* Fix failing tests in availability distribution

* Add traceID tag to launch-approval and validation-result

* Reshuffle the validation and validation result spans to where more appropriate and add block-hash tag

* - Add tranche and should-trigger tag to process-wakeup span
- Add candidate-hash and traceID to check-and-import-approval span

* cargo fmt

* - Adjustments after PR comments

* Move span pruning after other pruning logic

* Remove DerefMut - no longer needed

* Relabel request-chunk spans

* - Fix typo in span label
- Add docs for drops

* Add new approval-voting span pruning logic

* Undo removal of !

* cargo fmt

polkadot/node/network/approval-distribution/src/lib.rs

@@ -21,6 +21,7 @@
 #![warn(missing_docs)]
 
 use futures::{channel::oneshot, FutureExt as _};
+use polkadot_node_jaeger as jaeger;
 use polkadot_node_network_protocol::{
 	self as net_protocol,
 	grid_topology::{RandomRouting, RequiredRouting, SessionGridTopologies, SessionGridTopology},
@@ -35,7 +36,7 @@ use polkadot_node_subsystem::{
 		ApprovalCheckResult, ApprovalDistributionMessage, ApprovalVotingMessage,
 		AssignmentCheckResult, NetworkBridgeEvent, NetworkBridgeTxMessage,
 	},
-	overseer, ActiveLeavesUpdate, FromOrchestra, OverseerSignal, SpawnedSubsystem, SubsystemError,
+	overseer, FromOrchestra, OverseerSignal, SpawnedSubsystem, SubsystemError,
 };
 use polkadot_primitives::{
 	BlockNumber, CandidateIndex, Hash, SessionIndex, ValidatorIndex, ValidatorSignature,
@@ -180,6 +181,9 @@ struct State {
 
 	/// Config for aggression.
 	aggression_config: AggressionConfig,
+
+	/// HashMap from active leaves to spans
+	spans: HashMap<Hash, jaeger::PerLeafSpan>,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -390,9 +394,18 @@ impl State {
 	) {
 		let mut new_hashes = HashSet::new();
 		for meta in &metas {
+			let mut span = self
+				.spans
+				.get(&meta.hash)
+				.map(|span| span.child(&"handle-new-blocks"))
+				.unwrap_or_else(|| jaeger::Span::new(meta.hash, &"handle-new-blocks"))
+				.with_string_tag("block-hash", format!("{:?}", meta.hash))
+				.with_stage(jaeger::Stage::ApprovalDistribution);
+
 			match self.blocks.entry(meta.hash) {
 				hash_map::Entry::Vacant(entry) => {
 					let candidates_count = meta.candidates.len();
+					span.add_uint_tag("candidates-count", candidates_count as u64);
 					let mut candidates = Vec::with_capacity(candidates_count);
 					candidates.resize_with(candidates_count, Default::default);
 
@@ -690,6 +703,7 @@ impl State {
 			if let Some(block_entry) = self.blocks.remove(relay_block) {
 				self.topologies.dec_session_refs(block_entry.session);
 			}
+			self.spans.remove(&relay_block);
 		});
 
 		// If a block was finalized, this means we may need to move our aggression
@@ -1230,6 +1244,14 @@ impl State {
 	) -> HashMap<ValidatorIndex, ValidatorSignature> {
 		let mut all_sigs = HashMap::new();
 		for (hash, index) in indices {
+			let _span = self
+				.spans
+				.get(&hash)
+				.map(|span| span.child("get-approval-signatures"))
+				.unwrap_or_else(|| jaeger::Span::new(&hash, "get-approval-signatures"))
+				.with_string_tag("block-hash", format!("{:?}", hash))
+				.with_stage(jaeger::Stage::ApprovalDistribution);
+
 			let block_entry = match self.blocks.get(&hash) {
 				None => {
 					gum::debug!(
@@ -1650,13 +1672,18 @@ impl ApprovalDistribution {
 			match message {
 				FromOrchestra::Communication { msg } =>
 					Self::handle_incoming(&mut ctx, state, msg, &self.metrics, rng).await,
-				FromOrchestra::Signal(OverseerSignal::ActiveLeaves(ActiveLeavesUpdate {
-					..
-				})) => {
+				FromOrchestra::Signal(OverseerSignal::ActiveLeaves(update)) => {
 					gum::trace!(target: LOG_TARGET, "active leaves signal (ignored)");
 					// the relay chain blocks relevant to the approval subsystems
 					// are those that are available, but not finalized yet
-					// actived and deactivated heads hence are irrelevant to this subsystem
+					// actived and deactivated heads hence are irrelevant to this subsystem, other than
+					// for tracing purposes.
+					if let Some(activated) = update.activated {
+						let head = activated.hash;
+						let approval_distribution_span =
+							jaeger::PerLeafSpan::new(activated.span, "approval-distribution");
+						state.spans.insert(head, approval_distribution_span);
+					}
 				},
 				FromOrchestra::Signal(OverseerSignal::BlockFinalized(_hash, number)) => {
 					gum::trace!(target: LOG_TARGET, number = %number, "finalized signal");
@@ -1682,6 +1709,14 @@ impl ApprovalDistribution {
 				state.handle_new_blocks(ctx, metrics, metas, rng).await;
 			},
 			ApprovalDistributionMessage::DistributeAssignment(cert, candidate_index) => {
+				let _span = state
+					.spans
+					.get(&cert.block_hash)
+					.map(|span| span.child("import-and-distribute-assignment"))
+					.unwrap_or_else(|| jaeger::Span::new(&cert.block_hash, "distribute-assignment"))
+					.with_string_tag("block-hash", format!("{:?}", cert.block_hash))
+					.with_stage(jaeger::Stage::ApprovalDistribution);
+
 				gum::debug!(
 					target: LOG_TARGET,
 					"Distributing our assignment on candidate (block={}, index={})",
@@ -1701,6 +1736,14 @@ impl ApprovalDistribution {
 					.await;
 			},
 			ApprovalDistributionMessage::DistributeApproval(vote) => {
+				let _span = state
+					.spans
+					.get(&vote.block_hash)
+					.map(|span| span.child("import-and-distribute-approval"))
+					.unwrap_or_else(|| jaeger::Span::new(&vote.block_hash, "distribute-approval"))
+					.with_string_tag("block-hash", format!("{:?}", vote.block_hash))
+					.with_stage(jaeger::Stage::ApprovalDistribution);
+
 				gum::debug!(
 					target: LOG_TARGET,
 					"Distributing our approval vote on candidate (block={}, index={})",