Restores accept_span and uses in_scope() instead of enter()

ZcashFoundation · Jun 12, 2023 · 8c642c4 · 8c642c4
1 parent d18f814
commit 8c642c4
Showing 1 changed file with 35 additions and 28 deletions.
diff --git a/zebra-network/src/peer_set/initialize.rs b/zebra-network/src/peer_set/initialize.rs
@@ -603,39 +603,46 @@ where
             );
 
             let connected_addr = peer::ConnectedAddr::new_inbound_direct(addr);
-            info!(peer = ?connected_addr, "listen_accept");
-
-            debug!("got incoming connection");
-            handshaker.ready().await?;
-            // TODO: distinguish between proxied listeners and direct listeners
-            let handshaker_span = info_span!("listen_handshaker", peer = ?connected_addr);
-
-            // Construct a handshake future but do not drive it yet....
-            let handshake = handshaker.call(HandshakeRequest {
-                data_stream: tcp_stream,
-                connected_addr,
-                connection_tracker,
+            let accept_span = info_span!("listen_accept", peer = ?connected_addr);
+            accept_span.in_scope(|| {
+                debug!("got incoming connection");
             });
-            // ... instead, spawn a new task to handle this connection
-            {
-                let mut peerset_tx = peerset_tx.clone();
 
-                let handshake_task = tokio::spawn(
-                    async move {
-                        let handshake_result = handshake.await;
+            handshaker.ready().await?;
 
-                        if let Ok(client) = handshake_result {
-                            // The connection limit makes sure this send doesn't block
-                            let _ = peerset_tx.send((addr, client)).await;
-                        } else {
-                            debug!(?handshake_result, "error handshaking with inbound peer");
+            accept_span.in_scope(|| {
+                debug!("got incoming connection");
+
+                // TODO: distinguish between proxied listeners and direct listeners
+                let handshaker_span = info_span!("listen_handshaker", peer = ?connected_addr);
+
+                // Construct a handshake future but do not drive it yet....
+                let handshake = handshaker.call(HandshakeRequest {
+                    data_stream: tcp_stream,
+                    connected_addr,
+                    connection_tracker,
+                });
+                // ... instead, spawn a new task to handle this connection
+                {
+                    let mut peerset_tx = peerset_tx.clone();
+
+                    let handshake_task = tokio::spawn(
+                        async move {
+                            let handshake_result = handshake.await;
+
+                            if let Ok(client) = handshake_result {
+                                // The connection limit makes sure this send doesn't block
+                                let _ = peerset_tx.send((addr, client)).await;
+                            } else {
+                                debug!(?handshake_result, "error handshaking with inbound peer");
+                            }
                         }
-                    }
-                    .instrument(handshaker_span),
-                );
+                        .instrument(handshaker_span),
+                    );
 
-                handshakes.push(Box::pin(handshake_task));
-            }
+                    handshakes.push(Box::pin(handshake_task));
+                }
+            });
 
             // Rate-limit inbound connection handshakes.
             // But sleep longer after a successful connection,