fix: wip

Author: conblem

src/api/mod.rs

@@ -1,18 +1,18 @@
-use anyhow::{Result, Error};
-use futures_util::future::{OptionFuture, Future};
+use anyhow::Result;
+use futures_util::future::{Future, OptionFuture};
 use futures_util::stream::Stream;
-use futures_util::{FutureExt, StreamExt, TryStreamExt};
+use futures_util::{FutureExt, StreamExt};
 use hyper::server::conn::Http;
 use metrics::{metrics, metrics_wrapper};
 use sqlx::PgPool;
 use std::sync::Arc;
-use tokio::io::{AsyncRead, AsyncWrite, Result as IoResult};
+use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::TcpListener;
-use tracing::{info, error, info_span, Instrument};
+use tracing::{error, info, info_span, Instrument};
 use warp::{Filter, Rejection, Reply};
-use std::error::Error as StdError;
 
 use crate::config::Listener;
+use std::fmt::Display;
 
 mod metrics;
 mod proxy;
@@ -24,31 +24,40 @@ where
     I: Stream<Item = Result<S, E>> + Unpin + Send,
     S: Future<Output = Result<T, E>> + Send + 'static,
     T: AsyncRead + AsyncWrite + Send + Unpin + 'static,
-    E: Into<Error>,
+    E: Display,
     R: Filter<Error = Rejection> + Clone + Send + 'static,
     R::Extract: Reply,
 {
     let service = warp::service(routes);
     let http = Arc::new(Http::new());
 
     loop {
-        let span = info_span!("test");
+        let span = info_span!("TCP");
         let conn = match io.next().instrument(span.clone()).await {
             Some(Ok(conn)) => conn,
-            Some(Err(e)) => {
-                span.in_scope(|| error!("{}", e));
+            Some(Err(err)) => {
+                span.in_scope(|| error!("{}", err));
                 continue;
-            },
+            }
             None => break,
         };
 
         let http = http.clone();
         let service = service.clone();
 
-        tokio::spawn(async move {
-            let conn = conn.await.unwrap();
-            http.serve_connection(conn, service).await
-        }.instrument(span));
+        tokio::spawn(
+            async move {
+                let conn = match conn.await {
+                    Ok(conn) => conn,
+                    Err(err) => {
+                        error!("{}", err);
+                        return;
+                    }
+                };
+                http.serve_connection(conn, service).await;
+            }
+            .instrument(span),
+        );
     }
 }
 
@@ -68,18 +77,18 @@ pub async fn new(
 
     let http = http
         .map(move |http| proxy::wrap(http, http_proxy))
-        .map(|http| serve(http, routes.clone()))
+        .map(|http| serve(http, routes.clone()).instrument(info_span!("HTTP")))
         .map(tokio::spawn);
 
     let prom = prom
         .map(move |prom| proxy::wrap(prom, prom_proxy))
-        .map(|prom| serve(prom, metrics()))
+        .map(|prom| serve(prom, metrics()).instrument(info_span!("PROM")))
         .map(tokio::spawn);
 
     let https = https
         .map(move |https| proxy::wrap(https, https_proxy))
         .map(|https| tls::wrap(https, pool))
-        .map(|https| serve(https, routes))
+        .map(|https| serve(https, routes).instrument(info_span!("HTTPS")))
         .map(tokio::spawn);
 
     info!("Starting API");

src/api/proxy.rs

@@ -12,7 +12,7 @@ use tokio::net::{TcpListener, TcpStream};
 use tokio_stream::wrappers::TcpListenerStream;
 use tokio_util::io::poll_read_buf;
 use tracing::field::{display, Empty};
-use tracing::{debug_span, error, info, Instrument};
+use tracing::{debug_span, error, info, Instrument, Span};
 
 use crate::config::ProxyProtocol;
 
@@ -21,7 +21,11 @@ pub(super) fn wrap(
     proxy: ProxyProtocol,
 ) -> impl Stream<Item = IoResult<impl Future<Output = IoResult<ProxyStream>>>> + Send {
     TcpListenerStream::new(listener)
-        .map_ok(move |stream| stream.source(proxy))
+        .map_ok(move |stream| {
+            let span = Span::current();
+            span.record_all()
+            stream.source(proxy)
+        })
         .map_ok(|mut conn| {
             let span = debug_span!("ADDR", remote.addr = Empty);
             let span_clone = span.clone();

src/api/tls.rs

@@ -1,4 +1,4 @@
-use anyhow::{anyhow, Error, Result};
+use anyhow::{anyhow, Result};
 use futures_util::stream::{repeat, Stream};
 use futures_util::{StreamExt, TryFutureExt, TryStreamExt};
 use parking_lot::RwLock;
@@ -18,7 +18,11 @@ use crate::util::to_u64;
 pub(super) fn wrap<L, I>(
     listener: L,
     pool: PgPool,
-) -> impl Stream<Item = Result<impl Future<Output = Result<impl AsyncRead + AsyncWrite + Send + Unpin + 'static>>>> + Send
+) -> impl Stream<
+    Item = Result<
+        impl Future<Output = Result<impl AsyncRead + AsyncWrite + Send + Unpin + 'static>>,
+    >,
+> + Send
 where
     L: Stream<Item = IoResult<I>> + Send,
     I: Future<Output = IoResult<ProxyStream>> + Send,