fix: refactor

Author: conblem

src/api/mod.rs

@@ -1,9 +1,11 @@
-use anyhow::Result;
+use anyhow::{Error, Result};
 use futures_util::future::{Future, OptionFuture};
 use futures_util::stream::Stream;
-use futures_util::{FutureExt, StreamExt};
+use futures_util::{FutureExt, StreamExt, TryFutureExt};
 use hyper::server::conn::Http;
+use lazy_static::lazy_static;
 use metrics::{metrics, metrics_wrapper};
+use prometheus::{register_int_counter_vec, register_int_gauge_vec, IntCounterVec, IntGaugeVec};
 use sqlx::PgPool;
 use std::fmt::Display;
 use std::sync::Arc;
@@ -20,12 +22,27 @@ mod proxy;
 mod routes;
 mod tls;
 
-async fn serve<I, S, T, E, R>(mut io: I, routes: R)
+lazy_static! {
+    static ref TCP_TOTAL_CONNECTION_COUNTER: IntCounterVec = register_int_counter_vec!(
+        "tcp_total_connection_counter",
+        "Sum of TCP Connections",
+        &["endpoint"]
+    )
+    .unwrap();
+    static ref TCP_OPEN_CONNECTION_COUNTER: IntGaugeVec = register_int_gauge_vec!(
+        "tcp_open_connection_counter",
+        "Amount of currently open TCP Connections",
+        &["endpoint"]
+    )
+    .unwrap();
+}
+
+async fn serve<I, S, T, E, R>(mut io: I, routes: R, endpoint: &str)
 where
     I: Stream<Item = Result<S, E>> + Unpin + Send,
     S: Future<Output = Result<T, E>> + Send + 'static,
     T: AsyncRead + AsyncWrite + Send + Unpin + 'static,
-    E: Display,
+    E: Into<Error> + Display + Send,
     R: Filter<Error = Rejection> + Clone + Send + 'static,
     R::Extract: Reply,
 {
@@ -35,28 +52,28 @@ where
     loop {
         let span = info_span!("conn", remote.addr = Empty, remote.real = Empty);
         let conn = match io.next().instrument(span.clone()).await {
-            Some(Ok(conn)) => conn,
+            Some(Ok(conn)) => conn.err_into(),
             Some(Err(err)) => {
                 span.in_scope(|| error!("{}", err));
                 continue;
             }
             None => break,
         };
 
+        TCP_TOTAL_CONNECTION_COUNTER.with_label_values(&[endpoint]).inc();
+        let open_counter = TCP_OPEN_CONNECTION_COUNTER.with_label_values(&[endpoint]);
+        open_counter.inc();
+
         let http = http.clone();
         let service = service.clone();
 
         tokio::spawn(
             async move {
-                let conn = match conn.await {
-                    Ok(conn) => conn,
-                    Err(err) => {
-                        error!("{}", err);
-                        return;
-                    }
-                };
-                http.serve_connection(conn, service).await;
+                let conn = conn.await?;
+                Ok(http.serve_connection(conn, service).await?)
             }
+            .inspect_err(|err: &Error| error!("{}", err))
+            .inspect(move |_| open_counter.dec())
             .instrument(span),
         );
     }
@@ -78,18 +95,18 @@ pub async fn new(
 
     let http = http
         .map(move |http| proxy::wrap(http, http_proxy))
-        .map(|http| serve(http, routes.clone()).instrument(info_span!("HTTP")))
+        .map(|http| serve(http, routes.clone(), "HTTP").instrument(info_span!("HTTP")))
         .map(tokio::spawn);
 
     let prom = prom
         .map(move |prom| proxy::wrap(prom, prom_proxy))
-        .map(|prom| serve(prom, metrics()).instrument(info_span!("PROM")))
+        .map(|prom| serve(prom, metrics(), "PROM").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).instrument(info_span!("HTTPS")))
+        .map(|https| serve(https, routes, "HTTPS").instrument(info_span!("HTTPS")))
         .map(tokio::spawn);
 
     info!("Starting API");

src/api/proxy.rs

@@ -24,10 +24,9 @@ pub(super) fn wrap(
         .map_ok(move |stream| {
             let span = Span::current();
             span.record("remote.addr", &debug(stream.peer_addr()));
-            (stream.source(proxy), span)
-        })
-        .map_ok(|(mut conn, span)| {
             let span_clone = span.clone();
+
+            let mut conn = stream.source(proxy);
             async move {
                 match conn.proxy_peer().await {
                     Ok(Some(addr)) => {