use crate::{
connections::{Connections, Destination},
handshake,
message::Message,
metrics,
metrics::Metrics,
streams::mpsc,
};
#[cfg(test)]
pub(crate) mod mocks;
mod traits;
pub use traits::*;
use rand::RngCore;
use std::sync::Arc;
mod connections;
mod heartbeat;
pub use connections::{user_con, UserCon};
#[derive(Debug)]
pub(crate) enum ConnectError {
IO(std::io::Error),
Handshake(handshake::HandshakeError),
}
impl From<std::io::Error> for ConnectError {
fn from(other: std::io::Error) -> Self {
Self::IO(other)
}
}
impl From<handshake::HandshakeError> for ConnectError {
fn from(other: handshake::HandshakeError) -> Self {
Self::Handshake(other)
}
}
pub struct Client<M> {
server_destination: Destination,
external_port: u16,
key: Vec<u8>,
metrics: Arc<M>,
}
impl Client<metrics::Empty> {
pub fn new(server: Destination, external_port: u16, key: Vec<u8>) -> Self {
Self {
server_destination: server,
external_port,
key,
metrics: Arc::new(metrics::Empty::new()),
}
}
}
impl<M> Client<M> {
fn exponential_backoff(
attempt: u32,
max_time: Option<std::time::Duration>,
) -> std::time::Duration {
let raw_time = std::time::Duration::from_secs(2u64.pow(attempt));
let raw_jitter = rand::rngs::ThreadRng::default().next_u64() % 1000;
let raw_calced = raw_time.checked_add(std::time::Duration::from_millis(raw_jitter));
match (max_time, raw_calced) {
(Some(max), Some(calced)) if calced > max => max,
(_, Some(calced)) => calced,
(Some(max), _) => max,
_ => std::time::Duration::from_millis(0),
}
}
}
impl<M> Client<M>
where
M: Metrics + Send + Sync + 'static,
{
pub fn new_with_metrics(
server: Destination,
external_port: u16,
key: Vec<u8>,
metrics_collector: M,
) -> Self {
Self {
server_destination: server,
external_port,
key,
metrics: Arc::new(metrics_collector),
}
}
async fn start_con<H>(&self, handler: Arc<H>) -> Result<(), ConnectError>
where
H: Handler + Send + Sync + 'static,
{
info!("Establishing Connection...");
let target_addr = self.server_destination.get_full_address();
debug!("Conneting to server: {}", target_addr);
let mut connection = tokio::net::TcpStream::connect(target_addr).await?;
debug!("Connected to Server");
let handshake_conf = handshake::Config::new(self.external_port);
debug!("Starting Handshake...");
handshake::client::perform(&mut connection, &self.key, handshake_conf).await?;
debug!("Performed Handshake");
let (read_con, write_con) = connection.into_split();
info!("Established Conection");
let (queue_tx, queue_rx) = tokio::sync::mpsc::unbounded_channel();
let outgoing = std::sync::Arc::new(Connections::<mpsc::StreamWriter<Message>>::new());
tokio::task::spawn(heartbeat::keep_alive(
queue_tx.clone(),
std::time::Duration::from_secs(15),
));
tokio::task::spawn(connections::tx::sender(
write_con,
queue_rx,
self.metrics.clone(),
));
connections::rx::receiver(
read_con,
queue_tx.clone(),
outgoing,
handler,
self.metrics.clone(),
)
.await;
Ok(())
}
pub async fn start<H>(self, handler: Arc<H>) -> !
where
H: Handler + Send + Sync + 'static,
{
info!("Starting...");
let mut attempts = 0;
loop {
match self.start_con(handler.clone()).await {
Ok(_) => {
attempts = 0;
}
Err(e) => {
error!("Connecting: {:?}", e);
attempts += 1;
let wait_time = Self::exponential_backoff(
attempts,
Some(std::time::Duration::from_secs(60)),
);
info!("Waiting {:?} before trying to connect again", wait_time);
tokio::time::sleep(wait_time).await;
}
};
}
}
}