kiaproxy 0.1.2

High performance and extremely light TCP load balancer proxy.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

#![forbid(unsafe_code)]
use tokio::net::{TcpListener, TcpStream};
use tokio::time::{sleep, Duration};
use chrono::{SecondsFormat, Utc};
use uuid::Uuid;
use std::time::SystemTime;
use std::io;
use std::env;
use std::process;

async fn retry_connect(txid: &str, address: &str, max_retries: u32, delay: Duration) -> io::Result<TcpStream> {
    let mut retries = 0;

    while retries < max_retries {
        match TcpStream::connect(address).await {
            Ok(stream) => {
                return Ok(stream);
            }
            Err(e) => {
                retries += 1;
                let ts = chrono::DateTime::<Utc>::from(SystemTime::now()).to_rfc3339_opts(SecondsFormat::Millis, true);
                println!("{ts} - {txid} - ERROR: Failed to connect (attempt {} of {}) - {:?}", retries, max_retries, e);
                if retries < max_retries {
                    sleep(delay).await;
                } else {
                    return Err(e);
                }
            }
        }
    }

   Err(io::Error::other("max retries reached"))

}

async fn health(addr: &str) -> bool {
    TcpStream::connect(addr).await.is_ok()
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let listenerstr = match env::var("LISTENER") {
        Ok(val) => val,
        Err(_) => {
            println!("Set the environment variable LISTENER to the endpoint kiaproxy is to listen on.");
            println!("Example:");
            println!("  export LISTENER=0.0.0.0:443");
            process::exit(1);
        }

    };
    let srversstr = match env::var("SERVERS") {
        Ok(val) => val,
        Err(_) => {
            println!("Set the environment variable SERVERS to the ordered list of backends for kiaproxy to proxy and route to.");
            println!("Example:");
            println!("  export SERVERS=192.168.1.120:443,192.168.1.121:443,192.168.1.122:443");
            process::exit(1);
        }

    };

    let listener = TcpListener::bind(listenerstr).await?;
    let servers: Vec<String> = srversstr
        .split(',')
        .map(|s| s.trim().to_string())
        .collect();

    let strt = chrono::DateTime::<Utc>::from(SystemTime::now()).to_rfc3339_opts(SecondsFormat::Millis, true);
    println!("{strt} - INIT - INFO: kiaproxy v0.1.2 TCP load balancer listening on {:?} with backends {:?}", listener, servers);

    loop {
        let (mut inbound, addr) = listener.accept().await?;
        let servers = servers.clone();
        tokio::spawn(async move {
            let txid = Uuid::new_v4().to_string();
            let max_retries = 9;
            let delay = Duration::from_secs(1);

            let mut selected = servers[0].clone();

            for srv in servers {
                let ts = chrono::DateTime::<Utc>::from(SystemTime::now()).to_rfc3339_opts(SecondsFormat::Millis, true);
                println!("{ts} - {txid} - INFO: checking for backend {srv}");

                if health(&srv).await {
                    selected = srv.clone();
                    let ts = chrono::DateTime::<Utc>::from(SystemTime::now()).to_rfc3339_opts(SecondsFormat::Millis, true);
                    println!("{ts} - {txid} - INFO: selected first online backend {srv}");
                    break;
                }
            }

            match retry_connect(&txid, &selected, max_retries, delay).await {
                Ok(mut stream) => {
                    let ts = chrono::DateTime::<Utc>::from(SystemTime::now()).to_rfc3339_opts(SecondsFormat::Millis, true);
                    println!("{ts} - {txid} - INFO: {addr} connected to backend {:?}", stream.peer_addr());

                    if let Err(e) = tokio::io::copy_bidirectional(&mut inbound, &mut stream).await {
                        let ts = chrono::DateTime::<Utc>::from(SystemTime::now()).to_rfc3339_opts(SecondsFormat::Millis, true);
                        println!("{ts} - {txid} - INFO: connection drop {addr} - {e}");
                    }
                }
                Err(e) => {
                    let ts = chrono::DateTime::<Utc>::from(SystemTime::now()).to_rfc3339_opts(SecondsFormat::Millis, true);
                    println!("{ts} - {txid} - ERROR: Failed to connect after {} retries for {addr} - {:?}", max_retries, e);
                }
            }
        });
    }
}