From 9c675f9969bd1e3ba5977c563b866de0625a4553 Mon Sep 17 00:00:00 2001
From: katelyn martin <katelyn@penumbralabs.xyz>
Date: Thu, 21 Mar 2024 18:44:21 -0400
Subject: [PATCH] =?UTF-8?q?refactor(staking):=20=F0=9F=91=80=20attach=20a?=
 =?UTF-8?q?=20span=20to=20validator=20lookup=20futures?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

> This is used to model causal relationships such as when a single
> future spawns several related background tasks, et cetera.

\- https://docs.rs/tracing/latest/tracing/span/struct.Span.html#method.follows_from

NB: this commit is best viewed using `-w` to ignore whitespace changes.
---
 .../validator_handler/uptime_tracker.rs       | 57 ++++++++++++-------
 1 file changed, 35 insertions(+), 22 deletions(-)

diff --git a/crates/core/component/stake/src/component/validator_handler/uptime_tracker.rs b/crates/core/component/stake/src/component/validator_handler/uptime_tracker.rs
index 555fef19b3..c1803f56fc 100644
--- a/crates/core/component/stake/src/component/validator_handler/uptime_tracker.rs
+++ b/crates/core/component/stake/src/component/validator_handler/uptime_tracker.rs
@@ -13,7 +13,7 @@ use {
     std::collections::BTreeMap,
     tendermint::abci::types::CommitInfo,
     tokio::task::{AbortHandle, JoinSet},
-    tracing::instrument,
+    tracing::{error_span, instrument, Instrument},
 };
 
 /// A bundle of information about a validator used to track its uptime.
@@ -109,34 +109,47 @@ pub trait ValidatorUptimeTracker: StateWrite {
         identity_key: crate::IdentityKey,
         lookups: &mut Lookups,
     ) -> AbortHandle {
+        // Define, but do not yet `.await` upon, a collection of futures fetching information
+        // about a validator.
         let state = self.get_validator_state(&identity_key);
         let uptime = self.get_validator_uptime(&identity_key);
         let consensus_key = self.fetch_validator_consensus_key(&identity_key);
 
-        lookups.spawn(async move {
-            let state = state
-                .await?
-                .expect("every known validator must have a recorded state");
+        // Define a span indicating that the spawned future follows from the current context.
+        let span = {
+            let span = error_span!("fetching validator information", %identity_key);
+            let current = tracing::Span::current();
+            span.follows_from(current);
+            span
+        };
 
-            match state {
-                validator::State::Active => {
-                    // If the validator is active, we need its consensus key and current uptime data:
-                    Ok(Some((
-                        identity_key,
-                        consensus_key
-                            .await?
-                            .expect("every known validator must have a recorded consensus key"),
-                        uptime
-                            .await?
-                            .expect("every known validator must have a recorded uptime"),
-                    )))
-                }
-                _ => {
-                    // Otherwise, we don't need to track its uptime, and there's no data to fetch.
-                    Ok(None)
+        lookups.spawn(
+            async move {
+                let state = state
+                    .await?
+                    .expect("every known validator must have a recorded state");
+
+                match state {
+                    validator::State::Active => {
+                        // If the validator is active, we need its consensus key and current uptime data:
+                        Ok(Some((
+                            identity_key,
+                            consensus_key
+                                .await?
+                                .expect("every known validator must have a recorded consensus key"),
+                            uptime
+                                .await?
+                                .expect("every known validator must have a recorded uptime"),
+                        )))
+                    }
+                    _ => {
+                        // Otherwise, we don't need to track its uptime, and there's no data to fetch.
+                        Ok(None)
+                    }
                 }
             }
-        })
+            .instrument(span),
+        )
     }
 }