Fix propagating tracing spans into subchains

shotover-proxy/src/server.rs

@@ -183,75 +183,77 @@ impl<C: Codec + 'static> TcpCodecListener<C> {
  self.listener = Some(create_listener(&self.listen_addr).await?);
  }
  }
-
- // Accept a new socket. This will attempt to perform error handling.
- // The `accept` method internally attempts to recover errors, so an
- // error here is non-recoverable.
- let socket = self.accept().await?;
-
- debug!("got socket");
- self.available_connections_gauge
- .set(self.limit_connections.available_permits() as f64);
-
- let peer = socket
- .peer_addr()
- .map(|p| format!("{}", p.ip()))
- .unwrap_or_else(|_| "Unknown peer".to_string());
-
- let conn_string = socket
- .peer_addr()
- .map(|p| format!("{}:{}", p.ip(), p.port()))
- .unwrap_or_else(|_| "Unknown peer".to_string());
-
- // Create the necessary per-connection handler state.
- socket.set_nodelay(true)?;
-
- let (pushed_messages_tx, pushed_messages_rx) =
- tokio::sync::mpsc::unbounded_channel::<Messages>();
-
- let mut handler = Handler {
- chain: self
- .chain
- .clone_with_pushed_messages_tx(pushed_messages_tx.clone()),
- client_details: peer,
- conn_details: conn_string,
- source_details: self.source_name.clone(),
-
- // The connection state needs a handle to the max connections
- // semaphore. When the handler is done processing the
- // connection, a permit is added back to the semaphore.
- codec: self.codec.clone(),
- limit_connections: self.limit_connections.clone(),
-
- // Receive shutdown notifications.
- shutdown: Shutdown::new(self.trigger_shutdown_rx.clone()),
-
- terminate_tasks: None,
- tls: self.tls.clone(),
- timeout: self.timeout,
- };
-
  self.connection_count = self.connection_count.wrapping_add(1);
-
- // Spawn a new task to process the connections.
- tokio::spawn(
- async move {
- tracing::debug!("New connection from {}", handler.conn_details);
-
- // Process the connection. If an error is encountered, log it.
- if let Err(err) = handler.run(socket, pushed_messages_rx).await {
- error!(
- "{:?}",
- err.context("connection was unexpectedly terminated")
- );
- }
- }
- .instrument(tracing::error_span!(
- "connection",
- id = self.connection_count,
- source = self.source_name.as_str()
- )),
+ let span = tracing::error_span!(
+ "connection",
+ id = self.connection_count,
+ source = self.source_name.as_str(),
  );
+ async {
+ // Accept a new socket. This will attempt to perform error handling.
+ // The `accept` method internally attempts to recover errors, so an
+ // error here is non-recoverable.
+ let socket = self.accept().await?;
+
+ debug!("got socket");
+ self.available_connections_gauge
+ .set(self.limit_connections.available_permits() as f64);
+
+ let peer = socket
+ .peer_addr()
+ .map(|p| format!("{}", p.ip()))
+ .unwrap_or_else(|_| "Unknown peer".to_string());
+
+ let conn_string = socket
+ .peer_addr()
+ .map(|p| format!("{}:{}", p.ip(), p.port()))
+ .unwrap_or_else(|_| "Unknown peer".to_string());
+
+ // Create the necessary per-connection handler state.
+ socket.set_nodelay(true)?;
+
+ let (pushed_messages_tx, pushed_messages_rx) =
+ tokio::sync::mpsc::unbounded_channel::<Messages>();
+
+ let mut handler = Handler {
+ chain: self.chain.clone_with_pushed_messages_tx(pushed_messages_tx),
+ client_details: peer,
+ conn_details: conn_string,
+ source_details: self.source_name.clone(),
+
+ // The connection state needs a handle to the max connections
+ // semaphore. When the handler is done processing the
+ // connection, a permit is added back to the semaphore.
+ codec: self.codec.clone(),
+ limit_connections: self.limit_connections.clone(),
+
+ // Receive shutdown notifications.
+ shutdown: Shutdown::new(self.trigger_shutdown_rx.clone()),
+
+ terminate_tasks: None,
+ tls: self.tls.clone(),
+ timeout: self.timeout,
+ };
+
+ // Spawn a new task to process the connections.
+ tokio::spawn(
+ async move {
+ tracing::debug!("New connection from {}", handler.conn_details);
+
+ // Process the connection. If an error is encountered, log it.
+ if let Err(err) = handler.run(socket, pushed_messages_rx).await {
+ error!(
+ "{:?}",
+ err.context("connection was unexpectedly terminated")
+ );
+ }
+ }
+ .in_current_span(),
+ );
+ Ok::<(), anyhow::Error>(())
+ }
+ .instrument(span)
+ .await?;
  }
  }