mirror of https://github.com/chipsenkbeil/distant
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
360 lines
11 KiB
Rust
360 lines
11 KiB
Rust
3 years ago
|
mod handler;
|
||
|
mod state;
|
||
|
|
||
|
use state::State;
|
||
|
|
||
3 years ago
|
use crate::{
|
||
3 years ago
|
data::{Request, Response},
|
||
3 years ago
|
net::{Codec, DataStream, Transport, TransportListener, TransportReadHalf, TransportWriteHalf},
|
||
3 years ago
|
server::{
|
||
|
utils::{ConnTracker, ShutdownTask},
|
||
|
PortRange,
|
||
|
},
|
||
3 years ago
|
};
|
||
3 years ago
|
use futures::stream::{Stream, StreamExt};
|
||
3 years ago
|
use log::*;
|
||
3 years ago
|
use std::{net::IpAddr, sync::Arc};
|
||
3 years ago
|
use tokio::{
|
||
3 years ago
|
io::{self, AsyncRead, AsyncWrite},
|
||
|
net::TcpListener,
|
||
3 years ago
|
sync::{mpsc, Mutex},
|
||
3 years ago
|
task::{JoinError, JoinHandle},
|
||
|
time::Duration,
|
||
|
};
|
||
|
|
||
|
/// Represents a server that listens for requests, processes them, and sends responses
|
||
3 years ago
|
pub struct DistantServer {
|
||
3 years ago
|
conn_task: JoinHandle<()>,
|
||
|
}
|
||
|
|
||
3 years ago
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
||
|
pub struct DistantServerOptions {
|
||
|
pub shutdown_after: Option<Duration>,
|
||
|
pub max_msg_capacity: usize,
|
||
|
}
|
||
|
|
||
|
impl Default for DistantServerOptions {
|
||
|
fn default() -> Self {
|
||
|
Self {
|
||
|
shutdown_after: None,
|
||
|
max_msg_capacity: 1,
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
3 years ago
|
impl DistantServer {
|
||
|
/// Bind to an IP address and port from the given range, taking an optional shutdown duration
|
||
|
/// that will shutdown the server if there is no active connection after duration
|
||
3 years ago
|
pub async fn bind<U>(
|
||
3 years ago
|
addr: IpAddr,
|
||
|
port: PortRange,
|
||
3 years ago
|
codec: U,
|
||
3 years ago
|
opts: DistantServerOptions,
|
||
3 years ago
|
) -> io::Result<(Self, u16)>
|
||
|
where
|
||
|
U: Codec + Send + 'static,
|
||
|
{
|
||
3 years ago
|
debug!("Binding to {} in range {}", addr, port);
|
||
|
let listener = TcpListener::bind(port.make_socket_addrs(addr).as_slice()).await?;
|
||
|
|
||
|
let port = listener.local_addr()?.port();
|
||
|
debug!("Bound to port: {}", port);
|
||
|
|
||
3 years ago
|
let stream = TransportListener::initialize(listener, move |stream| {
|
||
|
Transport::new(stream, codec.clone())
|
||
|
})
|
||
3 years ago
|
.into_stream();
|
||
|
|
||
|
Ok((Self::initialize(Box::pin(stream), opts), port))
|
||
3 years ago
|
}
|
||
|
|
||
|
/// Initialize a distant server using the provided listener
|
||
3 years ago
|
pub fn initialize<T, U, S>(stream: S, opts: DistantServerOptions) -> Self
|
||
3 years ago
|
where
|
||
|
T: DataStream + Send + 'static,
|
||
3 years ago
|
U: Codec + Send + 'static,
|
||
|
S: Stream<Item = Transport<T, U>> + Send + Unpin + 'static,
|
||
3 years ago
|
{
|
||
3 years ago
|
// Build our state for the server
|
||
3 years ago
|
let state: Arc<Mutex<State>> = Arc::new(Mutex::new(State::default()));
|
||
3 years ago
|
let (shutdown, tracker) = ShutdownTask::maybe_initialize(opts.shutdown_after);
|
||
3 years ago
|
|
||
|
// Spawn our connection task
|
||
|
let conn_task = tokio::spawn(async move {
|
||
3 years ago
|
connection_loop(stream, state, tracker, shutdown, opts.max_msg_capacity).await
|
||
3 years ago
|
});
|
||
|
|
||
3 years ago
|
Self { conn_task }
|
||
3 years ago
|
}
|
||
|
|
||
|
/// Waits for the server to terminate
|
||
|
pub async fn wait(self) -> Result<(), JoinError> {
|
||
|
self.conn_task.await
|
||
|
}
|
||
3 years ago
|
|
||
|
/// Aborts the server by aborting the internal task handling new connections
|
||
|
pub fn abort(&self) {
|
||
|
self.conn_task.abort();
|
||
|
}
|
||
3 years ago
|
}
|
||
|
|
||
3 years ago
|
async fn connection_loop<T, U, S>(
|
||
3 years ago
|
mut stream: S,
|
||
3 years ago
|
state: Arc<Mutex<State>>,
|
||
|
tracker: Option<Arc<Mutex<ConnTracker>>>,
|
||
3 years ago
|
shutdown: Option<ShutdownTask>,
|
||
3 years ago
|
max_msg_capacity: usize,
|
||
3 years ago
|
) where
|
||
3 years ago
|
T: DataStream + Send + 'static,
|
||
3 years ago
|
U: Codec + Send + 'static,
|
||
|
S: Stream<Item = Transport<T, U>> + Send + Unpin + 'static,
|
||
3 years ago
|
{
|
||
3 years ago
|
let inner = async move {
|
||
|
loop {
|
||
3 years ago
|
match stream.next().await {
|
||
|
Some(transport) => {
|
||
3 years ago
|
let conn_id = rand::random();
|
||
3 years ago
|
debug!(
|
||
|
"<Conn @ {}> Established against {}",
|
||
|
conn_id,
|
||
3 years ago
|
transport.to_connection_tag()
|
||
3 years ago
|
);
|
||
3 years ago
|
if let Err(x) = on_new_conn(
|
||
3 years ago
|
transport,
|
||
3 years ago
|
conn_id,
|
||
3 years ago
|
Arc::clone(&state),
|
||
3 years ago
|
tracker.as_ref().map(Arc::clone),
|
||
|
max_msg_capacity,
|
||
|
)
|
||
|
.await
|
||
|
{
|
||
3 years ago
|
error!("<Conn @ {}> Failed handshake: {}", conn_id, x);
|
||
3 years ago
|
}
|
||
|
}
|
||
3 years ago
|
None => {
|
||
|
info!("Listener shutting down");
|
||
3 years ago
|
break;
|
||
|
}
|
||
3 years ago
|
};
|
||
3 years ago
|
}
|
||
3 years ago
|
};
|
||
|
|
||
3 years ago
|
match shutdown {
|
||
|
Some(shutdown) => tokio::select! {
|
||
|
_ = inner => {}
|
||
|
_ = shutdown => {
|
||
|
warn!("Reached shutdown timeout, so terminating");
|
||
|
}
|
||
|
},
|
||
|
None => inner.await,
|
||
3 years ago
|
}
|
||
|
}
|
||
|
|
||
|
/// Processes a new connection, performing a handshake, and then spawning two tasks to handle
|
||
|
/// input and output, returning join handles for the input and output tasks respectively
|
||
3 years ago
|
async fn on_new_conn<T, U>(
|
||
|
transport: Transport<T, U>,
|
||
3 years ago
|
conn_id: usize,
|
||
|
state: Arc<Mutex<State>>,
|
||
|
tracker: Option<Arc<Mutex<ConnTracker>>>,
|
||
3 years ago
|
max_msg_capacity: usize,
|
||
3 years ago
|
) -> io::Result<JoinHandle<()>>
|
||
3 years ago
|
where
|
||
|
T: DataStream,
|
||
3 years ago
|
U: Codec + Send + 'static,
|
||
3 years ago
|
{
|
||
3 years ago
|
// Update our tracker to reflect the new connection
|
||
|
if let Some(ct) = tracker.as_ref() {
|
||
|
ct.lock().await.increment();
|
||
|
}
|
||
|
|
||
3 years ago
|
// Split the transport into read and write halves so we can handle input
|
||
|
// and output concurrently
|
||
|
let (t_read, t_write) = transport.into_split();
|
||
|
let (tx, rx) = mpsc::channel(max_msg_capacity);
|
||
|
|
||
|
// Spawn a new task that loops to handle requests from the client
|
||
3 years ago
|
let state_2 = Arc::clone(&state);
|
||
|
let req_task = tokio::spawn(async move {
|
||
|
request_loop(conn_id, state_2, t_read, tx).await;
|
||
3 years ago
|
});
|
||
|
|
||
|
// Spawn a new task that loops to handle responses to the client
|
||
3 years ago
|
let res_task = tokio::spawn(async move { response_loop(conn_id, t_write, rx).await });
|
||
3 years ago
|
|
||
3 years ago
|
// Spawn cleanup task that waits on our req & res tasks to complete
|
||
|
let cleanup_task = tokio::spawn(async move {
|
||
|
// Wait for both receiving and sending tasks to complete before marking
|
||
|
// the connection as complete
|
||
|
let _ = tokio::join!(req_task, res_task);
|
||
|
|
||
|
state.lock().await.cleanup_connection(conn_id).await;
|
||
|
if let Some(ct) = tracker.as_ref() {
|
||
|
ct.lock().await.decrement();
|
||
|
}
|
||
|
});
|
||
|
|
||
|
Ok(cleanup_task)
|
||
3 years ago
|
}
|
||
|
|
||
|
/// Repeatedly reads in new requests, processes them, and sends their responses to the
|
||
|
/// response loop
|
||
3 years ago
|
async fn request_loop<T, U>(
|
||
3 years ago
|
conn_id: usize,
|
||
|
state: Arc<Mutex<State>>,
|
||
3 years ago
|
mut transport: TransportReadHalf<T, U>,
|
||
3 years ago
|
tx: mpsc::Sender<Response>,
|
||
3 years ago
|
) where
|
||
|
T: AsyncRead + Send + Unpin + 'static,
|
||
3 years ago
|
U: Codec,
|
||
3 years ago
|
{
|
||
3 years ago
|
loop {
|
||
|
match transport.receive::<Request>().await {
|
||
|
Ok(Some(req)) => {
|
||
|
debug!(
|
||
|
"<Conn @ {}> Received request of type{} {}",
|
||
3 years ago
|
conn_id,
|
||
3 years ago
|
if req.payload.len() > 1 { "s" } else { "" },
|
||
|
req.to_payload_type_string()
|
||
|
);
|
||
|
|
||
3 years ago
|
if let Err(x) = handler::process(conn_id, Arc::clone(&state), req, tx.clone()).await
|
||
|
{
|
||
|
error!("<Conn @ {}> {}", conn_id, x);
|
||
3 years ago
|
break;
|
||
|
}
|
||
|
}
|
||
|
Ok(None) => {
|
||
3 years ago
|
trace!("<Conn @ {}> Input from connection closed", conn_id);
|
||
3 years ago
|
break;
|
||
|
}
|
||
|
Err(x) => {
|
||
3 years ago
|
error!("<Conn @ {}> {}", conn_id, x);
|
||
3 years ago
|
break;
|
||
|
}
|
||
|
}
|
||
|
}
|
||
3 years ago
|
|
||
|
// Properly close off any associated process' stdin given that we can't get new
|
||
|
// requests to send more stdin to them
|
||
|
state.lock().await.close_stdin_for_connection(conn_id);
|
||
3 years ago
|
}
|
||
|
|
||
|
/// Repeatedly sends responses out over the wire
|
||
3 years ago
|
async fn response_loop<T, U>(
|
||
3 years ago
|
conn_id: usize,
|
||
3 years ago
|
mut transport: TransportWriteHalf<T, U>,
|
||
3 years ago
|
mut rx: mpsc::Receiver<Response>,
|
||
3 years ago
|
) where
|
||
|
T: AsyncWrite + Send + Unpin + 'static,
|
||
3 years ago
|
U: Codec,
|
||
3 years ago
|
{
|
||
3 years ago
|
while let Some(res) = rx.recv().await {
|
||
3 years ago
|
debug!(
|
||
|
"<Conn @ {}> Sending response of type{} {}",
|
||
|
conn_id,
|
||
|
if res.payload.len() > 1 { "s" } else { "" },
|
||
|
res.to_payload_type_string()
|
||
|
);
|
||
|
|
||
3 years ago
|
if let Err(x) = transport.send(res).await {
|
||
3 years ago
|
error!("<Conn @ {}> {}", conn_id, x);
|
||
3 years ago
|
break;
|
||
|
}
|
||
|
}
|
||
3 years ago
|
|
||
|
trace!("<Conn @ {}> Output to connection closed", conn_id);
|
||
3 years ago
|
}
|
||
3 years ago
|
|
||
|
#[cfg(test)]
|
||
|
mod tests {
|
||
|
use super::*;
|
||
3 years ago
|
use crate::{
|
||
|
data::{RequestData, ResponseData},
|
||
3 years ago
|
net::{InmemoryStream, PlainCodec},
|
||
3 years ago
|
};
|
||
|
use std::pin::Pin;
|
||
|
|
||
3 years ago
|
#[allow(clippy::type_complexity)]
|
||
3 years ago
|
fn make_transport_stream() -> (
|
||
3 years ago
|
mpsc::Sender<Transport<InmemoryStream, PlainCodec>>,
|
||
|
Pin<Box<dyn Stream<Item = Transport<InmemoryStream, PlainCodec>> + Send>>,
|
||
3 years ago
|
) {
|
||
3 years ago
|
let (tx, rx) = mpsc::channel::<Transport<InmemoryStream, PlainCodec>>(1);
|
||
3 years ago
|
let stream = futures::stream::unfold(rx, |mut rx| async move {
|
||
|
rx.recv().await.map(move |transport| (transport, rx))
|
||
|
});
|
||
|
(tx, Box::pin(stream))
|
||
3 years ago
|
}
|
||
|
|
||
3 years ago
|
#[tokio::test]
|
||
|
async fn wait_should_return_ok_when_all_inner_tasks_complete() {
|
||
|
let (tx, stream) = make_transport_stream();
|
||
|
|
||
|
let server = DistantServer::initialize(stream, Default::default());
|
||
3 years ago
|
|
||
3 years ago
|
// Conclude all server tasks by closing out the listener
|
||
|
drop(tx);
|
||
|
|
||
|
let result = server.wait().await;
|
||
|
assert!(result.is_ok(), "Unexpected result: {:?}", result);
|
||
3 years ago
|
}
|
||
|
|
||
3 years ago
|
#[tokio::test]
|
||
|
async fn wait_should_return_error_when_server_aborted() {
|
||
|
let (_tx, stream) = make_transport_stream();
|
||
|
|
||
|
let server = DistantServer::initialize(stream, Default::default());
|
||
|
server.abort();
|
||
|
|
||
|
match server.wait().await {
|
||
|
Err(x) if x.is_cancelled() => {}
|
||
|
x => panic!("Unexpected result: {:?}", x),
|
||
|
}
|
||
3 years ago
|
}
|
||
|
|
||
3 years ago
|
#[tokio::test]
|
||
|
async fn server_should_receive_requests_and_send_responses_to_appropriate_connections() {
|
||
|
let (tx, stream) = make_transport_stream();
|
||
|
|
||
|
let _server = DistantServer::initialize(stream, Default::default());
|
||
|
|
||
|
// Send over a "connection"
|
||
|
let (mut t1, t2) = Transport::make_pair();
|
||
|
tx.send(t2).await.unwrap();
|
||
|
|
||
|
// Send a request
|
||
|
t1.send(Request::new(
|
||
|
"test-tenant",
|
||
|
vec![RequestData::SystemInfo {}],
|
||
|
))
|
||
|
.await
|
||
|
.unwrap();
|
||
|
|
||
|
// Get a response
|
||
|
let res = t1.receive::<Response>().await.unwrap().unwrap();
|
||
|
assert!(res.payload.len() == 1, "Unexpected payload size");
|
||
|
assert!(
|
||
|
matches!(res.payload[0], ResponseData::SystemInfo { .. }),
|
||
|
"Unexpected response: {:?}",
|
||
|
res.payload[0]
|
||
|
);
|
||
3 years ago
|
}
|
||
|
|
||
3 years ago
|
#[tokio::test]
|
||
|
async fn server_should_shutdown_if_no_connections_after_shutdown_duration() {
|
||
|
let (_tx, stream) = make_transport_stream();
|
||
|
|
||
|
let server = DistantServer::initialize(
|
||
|
stream,
|
||
|
DistantServerOptions {
|
||
|
shutdown_after: Some(Duration::from_millis(50)),
|
||
|
max_msg_capacity: 1,
|
||
|
},
|
||
|
);
|
||
|
|
||
|
let result = server.wait().await;
|
||
|
assert!(result.is_ok(), "Unexpected result: {:?}", result);
|
||
3 years ago
|
}
|
||
|
}
|