From 6aea0fd9e8026e1985089418787cabfe5a7de904 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 23 Jun 2022 17:46:02 +1000
Subject: [PATCH] Add some missing tracing spans (#4660)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
---
 zebra-network/src/address_book_updater.rs   | 37 ++++++++++-----------
 zebra-network/src/peer_set/candidate_set.rs | 22 ++++++++----
 2 files changed, 33 insertions(+), 26 deletions(-)

diff --git a/zebra-network/src/address_book_updater.rs b/zebra-network/src/address_book_updater.rs
index dbb0eb16054..25e3d083a01 100644
--- a/zebra-network/src/address_book_updater.rs
+++ b/zebra-network/src/address_book_updater.rs
@@ -59,33 +59,32 @@ impl AddressBookUpdater {
         let address_book = Arc::new(std::sync::Mutex::new(address_book));
 
         let worker_address_book = address_book.clone();
-        let span = Span::current();
         let worker = move || {
-            span.in_scope(|| {
-                info!("starting the address book updater");
+            info!("starting the address book updater");
 
-                while let Some(event) = worker_rx.blocking_recv() {
-                    trace!(?event, "got address book change");
+            while let Some(event) = worker_rx.blocking_recv() {
+                trace!(?event, "got address book change");
 
-                    // # Correctness
-                    //
-                    // Briefly hold the address book threaded mutex, to update the
-                    // state for a single address.
-                    worker_address_book
-                        .lock()
-                        .expect("mutex should be unpoisoned")
-                        .update(event);
-                }
+                // # Correctness
+                //
+                // Briefly hold the address book threaded mutex, to update the
+                // state for a single address.
+                worker_address_book
+                    .lock()
+                    .expect("mutex should be unpoisoned")
+                    .update(event);
+            }
 
-                let error = Err(AllAddressBookUpdaterSendersClosed.into());
-                info!(?error, "stopping address book updater");
-                error
-            })
+            let error = Err(AllAddressBookUpdaterSendersClosed.into());
+            info!(?error, "stopping address book updater");
+            error
         };
 
         // Correctness: spawn address book accesses on a blocking thread,
         //              to avoid deadlocks (see #1976)
-        let address_book_updater_task_handle = tokio::task::spawn_blocking(worker);
+        let span = Span::current();
+        let address_book_updater_task_handle =
+            tokio::task::spawn_blocking(move || span.in_scope(worker));
 
         (
             address_book,
diff --git a/zebra-network/src/peer_set/candidate_set.rs b/zebra-network/src/peer_set/candidate_set.rs
index 950f05672a7..e072b37e716 100644
--- a/zebra-network/src/peer_set/candidate_set.rs
+++ b/zebra-network/src/peer_set/candidate_set.rs
@@ -4,6 +4,7 @@ use chrono::Utc;
 use futures::stream::{FuturesUnordered, StreamExt};
 use tokio::time::{sleep_until, timeout, Instant};
 use tower::{Service, ServiceExt};
+use tracing::Span;
 
 use zebra_chain::serialization::DateTime32;
 
@@ -333,9 +334,12 @@ where
         //
         // Extend handles duplicate addresses internally.
         let address_book = self.address_book.clone();
-        tokio::task::spawn_blocking(move || address_book.lock().unwrap().extend(addrs))
-            .await
-            .expect("panic in new peers address book update task");
+        let span = Span::current();
+        tokio::task::spawn_blocking(move || {
+            span.in_scope(|| address_book.lock().unwrap().extend(addrs))
+        })
+        .await
+        .expect("panic in new peers address book update task");
     }
 
     /// Returns the next candidate for a connection attempt, if any are available.
@@ -386,7 +390,8 @@ where
         };
 
         // Correctness: Spawn address book accesses on a blocking thread, to avoid deadlocks (see #1976).
-        let next_peer = tokio::task::spawn_blocking(next_peer)
+        let span = Span::current();
+        let next_peer = tokio::task::spawn_blocking(move || span.in_scope(next_peer))
             .await
             .expect("panic in next peer address book task")?;
 
@@ -406,9 +411,12 @@ where
         // Spawn address book accesses on a blocking thread,
         // to avoid deadlocks (see #1976).
         let address_book = self.address_book.clone();
-        tokio::task::spawn_blocking(move || address_book.lock().unwrap().update(addr))
-            .await
-            .expect("panic in peer failure address book update task");
+        let span = Span::current();
+        tokio::task::spawn_blocking(move || {
+            span.in_scope(|| address_book.lock().unwrap().update(addr))
+        })
+        .await
+        .expect("panic in peer failure address book update task");
     }
 }