feat: always ready for a streamer at relay

This commit is contained in:
Ahmet Kaan GÜMÜŞ 2024-04-25 04:19:03 +03:00
parent 5bdc68b6ee
commit 56d17084e1
3 changed files with 159 additions and 112 deletions

View file

@ -18,21 +18,13 @@ use tokio_rustls::{
rustls::pki_types::{CertificateDer, PrivateKeyDer},
TlsAcceptor,
};
use tokio_tungstenite::tungstenite::{Error, Message};
use tokio_tungstenite::tungstenite::{util::NonBlockingResult, Error, Message};
use crate::{Config, Listener, Streamer};
const BUFFER_LENGTH: usize = 1000000;
const MAX_TOLERATED_MESSAGE_COUNT: usize = 10;
pub async fn start(relay_configs: Config) {
//need to move them for multi streamer
let listener_socket = TcpListener::bind(relay_configs.listener_address)
.await
.unwrap();
let (record_producer, record_consumer) = channel(BUFFER_LENGTH);
let streamer_socket = TcpListener::bind(relay_configs.streamer_address)
.await
.unwrap();
let timer = Instant::now();
let fullchain: io::Result<Vec<CertificateDer<'static>>> = certs(&mut BufReader::new(
@ -53,114 +45,159 @@ pub async fn start(relay_configs: Config) {
.with_single_cert(fullchain, privkey)
.unwrap();
let acceptor = TlsAcceptor::from(Arc::new(server_tls_config));
let (streamer_alive_producer, streamer_alive_receiver) = tokio::sync::oneshot::channel();
let message_organizer_task: Option<JoinHandle<()>>;
let buffer_layer_task: Option<JoinHandle<()>>;
let (listener_stream_tasks_producer, listener_stream_tasks_receiver) =
tokio::sync::mpsc::channel(BUFFER_LENGTH);
let mut new_streamer = Streamer {
ip: "127.0.0.1".to_string().parse().unwrap(),
port: 0000,
};
loop {
match streamer_socket.accept().await {
Ok((streamer_tcp, streamer_info)) => {
new_streamer.ip = streamer_info.ip();
new_streamer.port = streamer_info.port();
println!(
"New Streamer: {:#?} | {:#?}",
streamer_info,
timer.elapsed()
);
if relay_configs.tls {
let streamer_tcp_tls = acceptor.accept(streamer_tcp).await.unwrap();
match tokio_tungstenite::accept_async(streamer_tcp_tls).await {
Ok(ws_stream) => {
tokio::spawn(streamer_stream(
new_streamer.clone(),
record_producer,
ws_stream,
timer,
streamer_alive_producer,
));
break;
//need to move them for multi streamer
let listener_socket = TcpListener::bind(relay_configs.listener_address.clone())
.await
.unwrap();
let (record_producer, record_consumer) = channel(BUFFER_LENGTH);
let streamer_socket = TcpListener::bind(relay_configs.streamer_address.clone())
.await
.unwrap();
let (streamer_alive_producer, streamer_alive_receiver) = tokio::sync::oneshot::channel();
let message_organizer_task: Option<JoinHandle<()>>;
let buffer_layer_task: Option<JoinHandle<()>>;
let (listener_stream_tasks_producer, listener_stream_tasks_receiver) =
tokio::sync::mpsc::channel(BUFFER_LENGTH);
let mut new_streamer = Streamer {
ip: "127.0.0.1".to_string().parse().unwrap(),
port: 0000,
};
loop {
match streamer_socket.accept().await {
Ok((streamer_tcp, streamer_info)) => {
new_streamer.ip = streamer_info.ip();
new_streamer.port = streamer_info.port();
println!(
"New Streamer: {:#?} | {:#?}",
streamer_info,
timer.elapsed()
);
if relay_configs.tls {
let streamer_tcp_tls = acceptor.accept(streamer_tcp).await.unwrap();
match tokio_tungstenite::accept_async(streamer_tcp_tls).await {
Ok(ws_stream) => {
tokio::spawn(streamer_stream(
new_streamer.clone(),
record_producer,
ws_stream,
timer,
streamer_alive_producer,
));
break;
}
Err(err_val) => eprintln!("Error: TCP to WS Transform | {}", err_val),
}
Err(err_val) => eprintln!("Error: TCP to WS Transform | {}", err_val),
}
} else {
match tokio_tungstenite::accept_async(streamer_tcp).await {
Ok(ws_stream) => {
tokio::spawn(streamer_stream(
new_streamer.clone(),
record_producer,
ws_stream,
timer,
streamer_alive_producer,
));
break;
} else {
match tokio_tungstenite::accept_async(streamer_tcp).await {
Ok(ws_stream) => {
tokio::spawn(streamer_stream(
new_streamer.clone(),
record_producer,
ws_stream,
timer,
streamer_alive_producer,
));
break;
}
Err(err_val) => eprintln!("Error: TCP to WS Transform | {}", err_val),
}
Err(err_val) => eprintln!("Error: TCP to WS Transform | {}", err_val),
}
}
Err(err_val) => eprintln!("Error: TCP Accept Connection | {}", err_val),
}
Err(err_val) => eprintln!("Error: TCP Accept Connection | {}", err_val),
}
let (message_producer, message_consumer) = channel(BUFFER_LENGTH);
let (buffered_producer, _) = channel(BUFFER_LENGTH);
message_organizer_task = tokio::spawn(message_organizer(
message_producer.clone(),
record_consumer,
relay_configs.latency,
))
.into();
buffer_layer_task = tokio::spawn(buffer_layer(
message_consumer,
buffered_producer.clone(),
relay_configs.latency,
))
.into();
let (listener_socket_killer_producer, listener_socket_killer_receiver) =
tokio::sync::oneshot::channel();
let listener_handler_task = tokio::spawn(listener_handler(
listener_socket,
acceptor.clone(),
relay_configs.tls,
buffered_producer.clone(),
listener_stream_tasks_producer,
timer,
listener_socket_killer_receiver,
));
status_checker(
buffered_producer.clone(),
timer,
new_streamer,
streamer_alive_receiver,
message_organizer_task,
buffer_layer_task,
listener_stream_tasks_receiver,
listener_handler_task,
listener_socket_killer_producer,
relay_configs.listener_address.clone(),
)
.await;
drop(streamer_socket);
}
}
async fn listener_handler(
listener_socket: TcpListener,
acceptor: TlsAcceptor,
is_tls: bool,
buffered_producer: Sender<Message>,
listener_stream_tasks_producer: tokio::sync::mpsc::Sender<JoinHandle<()>>,
timer: Instant,
mut listener_socket_killer_receiver: tokio::sync::oneshot::Receiver<bool>,
) {
while let Err(_) = listener_socket_killer_receiver.try_recv() {
match listener_socket.accept().await.no_block() {
Ok(accepted_request) => match accepted_request {
Some((tcp_stream, listener_info)) => {
let new_listener = Listener {
ip: listener_info.ip(),
port: listener_info.port(),
};
if is_tls {
let streamer_tcp_tls = acceptor.accept(tcp_stream).await.unwrap();
let wss_stream = tokio_tungstenite::accept_async(streamer_tcp_tls)
.await
.unwrap();
let listener_stream_task = tokio::spawn(stream(
new_listener,
wss_stream,
buffered_producer.subscribe(),
));
let _ = listener_stream_tasks_producer
.send(listener_stream_task)
.await;
} else {
let ws_stream = tokio_tungstenite::accept_async(tcp_stream).await.unwrap();
let listener_stream_task = tokio::spawn(stream(
new_listener,
ws_stream,
buffered_producer.subscribe(),
));
let _ = listener_stream_tasks_producer
.send(listener_stream_task)
.await;
}
println!("New Listener: {} | {:#?}", listener_info, timer.elapsed());
}
None => {}
},
Err(_) => {}
}
}
let (message_producer, message_consumer) = channel(BUFFER_LENGTH);
let (buffered_producer, _) = channel(BUFFER_LENGTH);
message_organizer_task = tokio::spawn(message_organizer(
message_producer.clone(),
record_consumer,
relay_configs.latency,
))
.into();
buffer_layer_task = tokio::spawn(buffer_layer(
message_consumer,
buffered_producer.clone(),
relay_configs.latency,
))
.into();
tokio::spawn(status_checker(
buffered_producer.clone(),
timer,
new_streamer,
streamer_alive_receiver,
message_organizer_task,
buffer_layer_task,
listener_stream_tasks_receiver,
));
while let Ok((tcp_stream, listener_info)) = listener_socket.accept().await {
let new_listener = Listener {
ip: listener_info.ip(),
port: listener_info.port(),
};
if relay_configs.tls {
let streamer_tcp_tls = acceptor.accept(tcp_stream).await.unwrap();
let wss_stream = tokio_tungstenite::accept_async(streamer_tcp_tls)
.await
.unwrap();
let listener_stream_task = tokio::spawn(stream(
new_listener,
wss_stream,
buffered_producer.subscribe(),
));
let _ = listener_stream_tasks_producer
.send(listener_stream_task)
.await;
} else {
let ws_stream = tokio_tungstenite::accept_async(tcp_stream).await.unwrap();
let listener_stream_task = tokio::spawn(stream(
new_listener,
ws_stream,
buffered_producer.subscribe(),
));
let _ = listener_stream_tasks_producer
.send(listener_stream_task)
.await;
}
println!("New Listener: {} | {:#?}", listener_info, timer.elapsed());
}
drop(listener_socket);
}
async fn status_checker(
buffered_producer: Sender<Message>,
@ -170,6 +207,9 @@ async fn status_checker(
message_organizer_task: Option<JoinHandle<()>>,
buffer_layer_task: Option<JoinHandle<()>>,
mut listener_stream_tasks_receiver: tokio::sync::mpsc::Receiver<JoinHandle<()>>,
listener_handler_task: JoinHandle<()>,
listener_socket_killer_producer: tokio::sync::oneshot::Sender<bool>,
listener_address: String,
) {
let mut listener_counter = buffered_producer.receiver_count();
let mut bottleneck_flag = false;
@ -185,6 +225,7 @@ async fn status_checker(
let cleaning_timer = Instant::now();
message_organizer_task.as_ref().unwrap().abort();
buffer_layer_task.as_ref().unwrap().abort();
listener_socket_killer_producer.send(true).unwrap();
let mut listener_task_counter = 0;
while listener_stream_tasks_receiver.len() > 0 {
match listener_stream_tasks_receiver.recv().await {
@ -195,6 +236,14 @@ async fn status_checker(
None => {}
}
}
if !listener_handler_task.is_finished() {
listener_handler_task.abort();
println!("Cleaning: Listener Handler Killed");
}
while TcpListener::bind(listener_address.clone()).await.is_err() {
tokio::time::sleep(Duration::from_millis(1)).await;
}
println!("Cleaning: Listener Socket Killed | {}", listener_address);
println!(
"Cleaning Done: Streamer Disconnected | {} | Disconnected Listener(s) = {} | {:#?}",
format!("{}:{}", streamer.ip, streamer.port),