From 13bbe2bee9570c6e7ae5f3d30c50ded4c6641bb1 Mon Sep 17 00:00:00 2001
From: Nadja Reitzenstein <me@dequbed.space>
Date: Fri, 24 Jun 2022 12:25:52 +0200
Subject: [PATCH] Add a connection-specific span to each API handler

---
 bffhd/capnp/connection.rs |  6 +++++-
 bffhd/capnp/mod.rs        | 40 +++++++++++++++++++++++++++++++++------
 bffhd/lib.rs              | 25 +++++++++++-------------
 3 files changed, 50 insertions(+), 21 deletions(-)

diff --git a/bffhd/capnp/connection.rs b/bffhd/capnp/connection.rs
index f8517e9..1f2a265 100644
--- a/bffhd/capnp/connection.rs
+++ b/bffhd/capnp/connection.rs
@@ -8,12 +8,14 @@ use crate::session::SessionManager;
 use capnp::capability::Promise;
 use capnp_rpc::pry;
 use rsasl::mechname::Mechname;
+use tracing::Span;
 
 /// Cap'n Proto API Handler
 pub struct BootCap {
     peer_addr: SocketAddr,
     authentication: AuthenticationHandle,
     sessionmanager: SessionManager,
+    span: Span,
 }
 
 impl BootCap {
@@ -21,12 +23,14 @@ impl BootCap {
         peer_addr: SocketAddr,
         authentication: AuthenticationHandle,
         sessionmanager: SessionManager,
+        span: Span,
     ) -> Self {
-        tracing::trace!(%peer_addr, "bootstrapping RPC");
+        tracing::trace!(parent: &span, %peer_addr, "bootstrapping RPC");
         Self {
             peer_addr,
             authentication,
             sessionmanager,
+            span,
         }
     }
 }
diff --git a/bffhd/capnp/mod.rs b/bffhd/capnp/mod.rs
index 42ba89f..e127ac0 100644
--- a/bffhd/capnp/mod.rs
+++ b/bffhd/capnp/mod.rs
@@ -12,7 +12,7 @@ use futures_util::{stream, AsyncRead, AsyncWrite, FutureExt, StreamExt};
 use std::future::Future;
 use std::io;
 
-use std::net::SocketAddr;
+use std::net::{IpAddr, SocketAddr};
 
 use crate::authentication::AuthenticationHandle;
 use crate::session::SessionManager;
@@ -145,12 +145,35 @@ impl APIServer {
         peer_addr: SocketAddr,
         stream: impl Future<Output = io::Result<TlsStream<IO>>>,
     ) {
-        tracing::debug!("handling new API connection");
+        let span = tracing::trace_span!("api.handle");
+        let _guard = span.enter();
+
+        struct Peer {
+            ip: IpAddr,
+            port: u16,
+        }
+
+        let peer = Peer {
+            ip: peer_addr.ip(),
+            port: peer_addr.port(),
+        };
+        tracing::debug!(
+            %peer.ip,
+            peer.port,
+            "spawning api handler"
+        );
+
+        let connection_span = tracing::info_span!(
+            "rpcsystem",
+            %peer.ip,
+            peer.port,
+        );
         let f = async move {
+            tracing::trace!(parent: &connection_span, "starting tls exchange");
             let stream = match stream.await {
                 Ok(stream) => stream,
-                Err(e) => {
-                    tracing::error!("TLS handshake failed: {}", e);
+                Err(error) => {
+                    tracing::error!(parent: &connection_span, %error, "TLS handshake failed");
                     return;
                 }
             };
@@ -161,10 +184,15 @@ impl APIServer {
                 peer_addr,
                 self.authentication.clone(),
                 self.sessionmanager.clone(),
+                connection_span.clone(),
             ));
 
-            if let Err(e) = RpcSystem::new(Box::new(vat), Some(bootstrap.client)).await {
-                tracing::error!("Error during RPC handling: {}", e);
+            if let Err(error) = RpcSystem::new(Box::new(vat), Some(bootstrap.client)).await {
+                tracing::error!(
+                    parent: &connection_span,
+                    %error,
+                    "error occured during rpc handling",
+                );
             }
         };
         let cgroup = SupervisionRegistry::with(SupervisionRegistry::new_group);
diff --git a/bffhd/lib.rs b/bffhd/lib.rs
index 777b81e..ec64ef5 100644
--- a/bffhd/lib.rs
+++ b/bffhd/lib.rs
@@ -43,7 +43,7 @@ mod tls;
 
 use std::sync::Arc;
 
-use futures_util::StreamExt;
+use futures_util::{FutureExt, StreamExt};
 use miette::{Context, IntoDiagnostic, Report};
 use once_cell::sync::OnceCell;
 
@@ -94,20 +94,17 @@ impl Diflouroborane {
         }
         tracing::info!("Server is being spawned");
         let handle = executor.spawn(server.serve());
-        std::thread::spawn(move || {
-            let result = async_io::block_on(handle);
-            match result {
-                Some(Ok(())) => {
-                    tracing::info!("console server finished without error");
-                }
-                Some(Err(error)) => {
-                    tracing::info!(%error, "console server finished with error");
-                }
-                None => {
-                    tracing::info!("console server finished with panic");
-                }
+        executor.spawn(handle.map(|result| match result {
+            Some(Ok(())) => {
+                tracing::info!("console server finished without error");
             }
-        });
+            Some(Err(error)) => {
+                tracing::info!(%error, "console server finished with error");
+            }
+            None => {
+                tracing::info!("console server finished with panic");
+            }
+        }));
 
         let env = StateDB::open_env(&config.db_path)?;