feat: handle SIGTREM

src/proxy.rs

@@ -1,11 +1,12 @@
 use evalexpr::*;
+use std::collections::HashMap;
 use std::env;
 use std::net::SocketAddr;
 use std::time::Duration;
+use tokio::task::JoinHandle;
 
 use tokio::net::{TcpListener, TcpStream};
 use tracing::Instrument;
-use tracing_subscriber::{prelude::*, EnvFilter};
 
 use crate::mc_server::{self, MinecraftAPI, MinecraftServerHandle, ServerDeploymentStatus};
 use crate::opaque_error::OpaqueError;
@@ -18,15 +19,85 @@ use crate::Config;
 static BYE_MESSAGE: &str = concat!("§dTami§r with §d<3§r §8(rev: ", env!("COMMIT_HASH"), ")§r");
 static OFFLINE_TIMER: std::time::Duration = Duration::from_secs(600);
 
-pub async fn process_connection<T: MinecraftServerHandle>(
+pub fn start_proxy<T: MinecraftServerHandle>(
+    listener: TcpListener,
+    api: impl MinecraftAPI<T> + Clone,
+    config: Config,
+    token: tokio_util::sync::CancellationToken,
+) -> JoinHandle<()> {
+    tokio::spawn(async move {
+        let tracker = tokio_util::task::TaskTracker::new();
+        loop {
+            tokio::select! {
+                Ok((socket, addr)) = listener.accept() => {
+                    let api = api.clone();
+
+                    let config = config.clone();
+                    tracker.spawn(async move {
+                        tracing::debug!(
+                            addr = format!("{}:{}", addr.ip().to_string(), addr.port().to_string()),
+                            "Client connected"
+                        );
+                        if let Err(e) = process_connection(socket, addr, api, config).await {
+                            trace_opaque(&e, "Client disconnected");
+                        } else {
+                            tracing::debug!(
+                                addr = format!("{}:{}", addr.ip().to_string(), addr.port().to_string()),
+                                "Client disconnected"
+                            );
+                        }
+                    });
+                }
+                _ = token.cancelled() => {
+                    tracker.close();
+                    let open_connections = tracker.len();
+                    tracing::info!(open_connections,"stopped handling new connections");
+                    tracker.wait().await;
+                    break;
+                }
+            }
+        }
+    })
+}
+
+fn trace_opaque(e: &OpaqueError, str: &str) {
+    match e.level {
+        tracing::Level::ERROR => tracing::error!(
+            // addr = format!("{}:{}", addr.ip().to_string(), addr.port().to_string()),
+            trace = %e.print_span_trace(),
+            err = format!("{}", e.context),
+            message = str
+        ),
+        tracing::Level::WARN => tracing::warn!(
+            // addr = format!("{}:{}", addr.ip().to_string(), addr.port().to_string()),
+            trace = %e.print_span_trace(),
+            err = format!("{}", e.context),
+            message = str
+        ),
+        tracing::Level::INFO => tracing::info!(
+            // addr = format!("{}:{}", addr.ip().to_string(), addr.port().to_string()),
+            trace = %e.print_span_trace(),
+            err = format!("{}", e.context),
+            message = str
+        ),
+        _ => {
+            tracing::error!(
+                // addr = format!("{}:{}", addr.ip().to_string(), addr.port().to_string()),
+                trace = %e.print_span_trace(),
+                err = format!("{}", e.context),
+                actual_level = ?e.level,
+                "Client disconnected (bad level)"
+            )
+        }
+    }
+}
+
+async fn process_connection<T: MinecraftServerHandle>(
     mut client_stream: TcpStream,
     addr: SocketAddr,
-    api: impl MinecraftAPI<T> + Send + Sync + 'static + Clone,
+    api: impl MinecraftAPI<T> + Clone,
     config: Config,
-) -> Result<(), OpaqueError>
-where
-    T: Send + Sync + 'static,
-{
+) -> Result<(), OpaqueError> {
     // this is wrapper so that async doesnt mess up the span, and
     // to make sure this doesn't propagate to later `handle_*`
     #[tracing::instrument(level = "info", skip(client_stream, config))]
@@ -109,10 +180,7 @@ async fn handle_status<T: MinecraftServerHandle>(
     client_stream: &mut TcpStream,
     handshake: &Handshake,
     api: impl MinecraftAPI<T>,
-) -> Result<(), OpaqueError>
-where
-    T: Send + Sync + 'static,
-{
+) -> Result<(), OpaqueError> {
     let client_packet = Packet::parse(client_stream).await?;
     if client_packet.id.get_int() != 0 {
         return Err(OpaqueError::create(&format!(