feat: (tcp) proxy

Cargo.lock

@@ -279,7 +279,6 @@ dependencies = [
  "futures-util",
  "glob",
  "hex",
- "human-units",
  "log",
  "lz4",
  "nix",
@@ -587,15 +586,6 @@ version = "1.10.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87"
 
-[[package]]
-name = "human-units"
-version = "0.5.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "47cf34dbcbbb7f1f6589c18e26f15a5a72592750dd5472037eb78fc0f92020d4"
-dependencies = [
- "paste",
-]
-
 [[package]]
 name = "hyper"
 version = "1.8.1"
@@ -1069,12 +1059,6 @@ dependencies = [
  "winapi",
 ]
 
-[[package]]
-name = "paste"
-version = "1.0.15"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a"
-
 [[package]]
 name = "percent-encoding"
 version = "2.3.2"

Cargo.toml

@@ -24,7 +24,6 @@ futures = "0.3.31"
 futures-util = "0.3.31"
 glob = "0.3.2"
 hex = "0.4.3"
-human-units = "0.5.3"
 log = "0.4.27"
 lz4 = "1.28.1"
 nix = { version = "0.30.1", features = ["user"] }

src/bin/dkl.rs

@@ -1,8 +1,8 @@
 use clap::{CommandFactory, Parser, Subcommand};
 use eyre::{format_err, Result};
-use human_units::Duration;
 use log::{debug, error};
 use std::net::SocketAddr;
+use std::time::Duration;
 use tokio::fs;
 
 #[derive(Parser)]
@@ -71,12 +71,12 @@ enum Command {
         #[arg(long, short = 'l')]
         listen: Vec<SocketAddr>,
         targets: Vec<SocketAddr>,
-        /// target polling interval
+        /// target polling interval, in seconds
-        #[arg(long, default_value = "30s")]
+        #[arg(long, default_value = "30")]
-        poll: Duration,
+        poll: u16,
-        /// connect or check timeout
+        /// connect or check timeout, in seconds
-        #[arg(long, default_value = "5s")]
+        #[arg(long, default_value = "5")]
-        timeout: Duration,
+        timeout: u16,
     },
 }
 
@@ -147,8 +147,8 @@ async fn main() -> Result<()> {
         } => Ok(dkl::proxy::Proxy {
             listen_addrs: listen,
             targets,
-            poll: poll.into(),
+            poll_interval: Duration::from_secs(poll.into()),
-            timeout: timeout.into(),
+            timeout: Duration::from_secs(timeout.into()),
         }
         .run()
         .await

src/proxy.rs

@@ -1,17 +1,16 @@
-use log::{info, log_enabled, warn};
+use log::info;
 use std::convert::Infallible;
-use std::io;
 use std::net::SocketAddr;
-use std::sync::atomic::{AtomicBool, Ordering::Relaxed};
+use std::sync::atomic::Ordering::Relaxed;
+use std::sync::{atomic::AtomicBool, Arc};
 use std::time::Duration;
 use thiserror::Error;
 use tokio::net::{TcpListener, TcpStream};
-use tokio::time;
 
 pub struct Proxy {
     pub listen_addrs: Vec<SocketAddr>,
     pub targets: Vec<SocketAddr>,
-    pub poll: Duration,
+    pub poll_interval: Duration,
     pub timeout: Duration,
 }
 
@@ -37,96 +36,76 @@ impl Proxy {
 
         // all targets are initially ok (better land on a down one than just fail)
         let targets: Vec<_> = (self.targets.into_iter())
-            .map(|addr| TargetStatus {
+            .map(|t| TargetStatus(t, AtomicBool::new(true)))
-                addr,
-                up: AtomicBool::new(true),
-                timeout: self.timeout,
-            })
             .collect();
-
+        let targets = Arc::new(targets);
-        // the proxy runs forever -> using 'static is not a leak
-        let targets = targets.leak();
 
         for listener in listeners {
-            tokio::spawn(proxy_listener(listener, targets));
+            tokio::spawn(proxy_listener(listener, targets.clone()));
         }
 
-        check_targets(targets, self.poll).await
+        check_targets(targets.clone(), self.poll_interval, self.timeout).await
     }
 }
 
-struct TargetStatus {
+struct TargetStatus(SocketAddr, AtomicBool);
-    addr: SocketAddr,
-    up: AtomicBool,
-    timeout: Duration,
-}
 impl TargetStatus {
     fn is_up(&self) -> bool {
-        self.up.load(Relaxed)
+        self.1.load(Relaxed)
-    }
-
-    fn set_up(&self, is_up: bool) {
-        let prev = self.up.swap(is_up, Relaxed);
-        if prev != is_up {
-            if is_up {
-                info!("{} is up", self.addr);
-            } else {
-                warn!("{} is down", self.addr);
-            }
-        }
-    }
-
-    async fn connect(&self) -> io::Result<TcpStream> {
-        let r = match time::timeout(self.timeout, TcpStream::connect(self.addr)).await {
-            Ok(r) => r,
-            Err(e) => Err(io::Error::new(io::ErrorKind::TimedOut, e)),
-        };
-
-        self.set_up(r.is_ok());
-        r
     }
 }
 
-async fn check_targets(targets: &'static [TargetStatus], poll: Duration) -> ! {
+async fn check_targets(targets: Arc<Vec<TargetStatus>>, poll: Duration, timeout: Duration) -> ! {
     use tokio::time;
     let mut poll_ticker = time::interval(poll);
     poll_ticker.set_missed_tick_behavior(time::MissedTickBehavior::Skip);
 
+    let mut first = true;
+
     loop {
         poll_ticker.tick().await;
 
-        let mut tasks = tokio::task::JoinSet::new();
+        for i in 0..targets.len() {
-
+            let targets = targets.clone();
-        for target in targets {
+            tokio::spawn(async move { check_target(&targets[i], timeout, first).await });
-            tasks.spawn(target.connect());
         }
 
-        tasks.join_all().await;
+        first = false;
-
-        if log_enabled!(log::Level::Info) {
-            let mut infos = String::new();
-            for ts in targets.iter() {
-                infos.push_str(&format!("{} ", ts.addr));
-                infos.push_str(if ts.is_up() { "up   " } else { "down " });
     }
-            info!("{infos}");
+}
+
+async fn check_target(ts: &TargetStatus, timeout: Duration, first: bool) {
+    let target = ts.0;
+
+    let is_up = match tokio::time::timeout(timeout, TcpStream::connect(target)).await {
+        Ok(Ok(_)) => true,
+        _ => false,
+    };
+
+    let prev = ts.1.swap(is_up, Relaxed);
+
+    if first || prev != is_up {
+        if is_up {
+            info!("{target} is up");
+        } else {
+            info!("{target} is down");
         }
     }
 }
 
-async fn proxy_listener(listener: TcpListener, targets: &'static [TargetStatus]) {
+async fn proxy_listener(listener: TcpListener, targets: Arc<Vec<TargetStatus>>) {
     let mut rng = fastrand::Rng::new();
 
     loop {
         let mut active = Vec::with_capacity(targets.len());
         let (mut src, _) = listener.accept().await.expect("listener.accept() failed");
 
-        active.extend((targets.iter().enumerate()).filter_map(|(i, ts)| ts.is_up().then_some(i)));
+        active.extend(targets.iter().filter(|ts| ts.is_up()).map(|ts| ts.0));
         rng.shuffle(&mut active);
 
         tokio::spawn(async move {
-            for i in active {
+            for target in active {
-                if let Ok(mut dst) = targets[i].connect().await {
+                if let Ok(mut dst) = TcpStream::connect(target).await {
                     let _ = tokio::io::copy_bidirectional(&mut src, &mut dst).await;
                     break;
                 }