joeyahines
/
RaaS
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refactored, handle failures more gracefully 

+ Split ConnectionHandler into its own struct
+ Handle errors more gracefully
+ System no longer locks up or crashes if a bot goes offline
+ Bots who go offline are automatically removed from the connector pool
main
Joey Hines 2024-06-16 16:37:28 -06:00
parent c6031f5fa4
commit 0136772c23
No known key found for this signature in database
GPG Key ID: 995E531F7A569DDB
5 changed files with 205 additions and 140 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

45
src/main.rs
View File

@ -1,12 +1,17 @@
use crate::robot::connection_handler::ConnectionHandler;
use crate::robot::robot_manager::RobotManager;
use log::info;
use crate::robot::ConnectorManager;
use log::{error, info};
use prost::Message;
use raas_types::raas;
use raas_types::raas::cmd::Request;
use raas_types::raas::ping::Ping;
use raas_types::raas::{recv_raas_msg, send_raas_msg};
use std::sync::Arc;
use std::time::Duration;
use tokio::io::duplex;
use tokio::net::TcpListener;
use tokio::sync::Mutex;
mod robot;
@ -15,10 +20,25 @@ async fn main() {
env_logger::init();
info!("Starting...");
let (mut main_stream, manager_stream) = duplex(10);
let mut manager = RobotManager::new("0.0.0.0:8080", manager_stream);
let addr = "0.0.0.0:8080";
manager.start_manager().await.unwrap();
let connector_manager = ConnectorManager {
next_id: 0,
connectors: Default::default(),
};
let connector_manager = Arc::new(Mutex::new(connector_manager));
let tcp_listener = TcpListener::bind(addr).await.unwrap();
let (mut main_stream, manager_stream) = duplex(10);
let robot_manager = RobotManager::new(manager_stream);
let connector_handler = ConnectionHandler::new(addr, tcp_listener);
let robot_manager_connector_manager = connector_manager.clone();
tokio::task::spawn(robot_manager.worker(robot_manager_connector_manager));
tokio::task::spawn(connector_handler.worker(connector_manager));
loop {
info!("Sleeping");
@ -43,8 +63,21 @@ async fn main() {
send_raas_msg(&mut main_stream, msg).await.unwrap();
info!("Wait for resp");
let resp = recv_raas_msg(&mut main_stream).await.unwrap();
let ret =
tokio::time::timeout(Duration::from_secs(5), recv_raas_msg(&mut main_stream)).await;
info!("Got resp {:?}", resp);
match ret {
Ok(resp) => match resp {
Ok(resp) => {
info!("Got response {:?}", resp);
}
Err(err) => {
error!("Error getting response: {}", err)
}
},
Err(_) => {
error!("Timed out waiting for response!")
}
}
}
}

94
src/robot/connection_handler.rs 100644
View File

@ -0,0 +1,94 @@
use crate::robot::robot_connector::RobotConnector;
use crate::robot::ROBOT_MESSAGE_QUEUE_SIZE;
use crate::robot::{ConnectorManager, Error};
use log::{error, info};
use prost::Message;
use raas_types::raas::register::{Register, RegisterResponse};
use raas_types::raas::{recv_raas_msg, send_raas_msg};
use std::sync::Arc;
use tokio::io::duplex;
use tokio::net::TcpListener;
use tokio::sync::Mutex;
pub struct ConnectionHandler {
tcp_listener: TcpListener,
addr: String,
}
impl ConnectionHandler {
pub fn new(addr: &str, tcp_listener: TcpListener) -> Self {
Self {
tcp_listener,
addr: addr.to_string(),
}
}
async fn handle_register(
&mut self,
connector_manager: Arc<Mutex<ConnectorManager>>,
) -> Result<(), Error> {
info!("Listening for connection at {}", self.addr);
let (mut socket, addr) = self.tcp_listener.accept().await?;
info!("Got connection from {}", addr);
let register_msg = recv_raas_msg(&mut socket).await?;
let register = Register::decode(&*register_msg.msg)?;
let mut connector_manager = connector_manager.lock().await;
let id = connector_manager.next_id;
connector_manager.next_id += 1;
let register_resp = RegisterResponse {
name: register.name.to_string(),
r#type: register.bot_type,
id,
};
let mut msg = Vec::new();
register_resp.encode(&mut msg)?;
send_raas_msg(&mut socket, msg).await?;
let (connector_duplex_manager, connector_duplex_connector) =
duplex(ROBOT_MESSAGE_QUEUE_SIZE);
let mut connector = RobotConnector::new(
id,
vec![register.bot_type()],
socket,
connector_duplex_connector,
);
let connector_handler = crate::robot::ConnectorHandle {
stream: connector_duplex_manager,
tags: connector.bot_tags.clone(),
};
connector_manager.connectors.insert(id, connector_handler);
tokio::spawn(async move { connector.worker().await });
Ok(())
}
pub async fn worker(
mut self,
connector_manager: Arc<Mutex<ConnectorManager>>,
) -> Result<(), Error> {
loop {
let res = self.handle_register(connector_manager.clone()).await;
if let Err(err) = res {
error!("Got error handling new register: {}", err);
if let Error::Io(_io_err) = &err {
error!("IO Error, exiting...");
return Err(err);
}
}
}
}
}

30
src/robot/mod.rs
View File

@ -1,4 +1,34 @@
use raas_types::raas;
use std::collections::HashMap;
use thiserror::Error;
use tokio::io::DuplexStream;
pub mod connection_handler;
pub mod robot_connector;
pub mod robot_manager;
pub const ROBOT_MESSAGE_QUEUE_SIZE: usize = 10;
#[derive(Debug)]
pub struct ConnectorHandle {
pub stream: DuplexStream,
pub tags: Vec<raas::register::BotTypes>,
}
#[derive(Debug)]
pub struct ConnectorManager {
pub next_id: u32,
pub connectors: HashMap<u32, ConnectorHandle>,
}
#[derive(Error, Debug)]
pub enum Error {
#[error("IO Error")]
Io(#[from] std::io::Error),
#[error("Protobuf Encode Error")]
ProtobufEncode(#[from] prost::EncodeError),
#[error("Protobuf Decode Error")]
ProtobufDecode(#[from] prost::DecodeError),
#[error("Connection to bot has been closed")]
ConnectionClosed(u32),
}

16
src/robot/robot_connector.rs
View File

@ -1,23 +1,13 @@
use crate::robot::Error;
use log::{error, info};
use prost::Message;
use raas_types::raas::cmd::{Command, Request};
use raas_types::raas::register::BotTypes;
use raas_types::raas::resp::Response;
use raas_types::raas::{recv_raas_msg, send_raas_msg};
use thiserror::Error;
use tokio::io::DuplexStream;
use tokio::net::TcpStream;
#[derive(Error, Debug)]
pub enum Error {
#[error("IO Error")]
Io(#[from] std::io::Error),
#[error("Protobuf Encode Error")]
ProtobufEncode(#[from] prost::EncodeError),
#[error("Protobuf Decode Error")]
ProtobufDecode(#[from] prost::DecodeError),
}
#[derive(Debug)]
pub struct RobotConnector {
pub bot_id: u32,
@ -85,10 +75,6 @@ impl RobotConnector {
Ok(())
}
pub fn is_running(&self) -> bool {
self.is_running
}
pub async fn worker(&mut self) {
self.is_running = true;

160
src/robot/robot_manager.rs
View File

@ -1,129 +1,42 @@
use crate::robot::robot_connector::{Error, RobotConnector};
use crate::robot::ROBOT_MESSAGE_QUEUE_SIZE;
use crate::robot::{ConnectorManager, Error};
use log::{error, info};
use prost::Message;
use raas_types::raas;
use raas_types::raas::bot::roll::Roll;
use raas_types::raas::cmd::request::Cmd;
use raas_types::raas::cmd::{Command, Request};
use raas_types::raas::register::{Register, RegisterResponse};
use raas_types::raas::resp::response::Resp;
use raas_types::raas::resp::Response;
use raas_types::raas::cmd::Request;
use raas_types::raas::{recv_raas_msg, send_raas_msg, RaasMessage};
use std::collections::HashMap;
use std::ops::{Deref, DerefMut};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::io::{duplex, AsyncReadExt, AsyncWriteExt, DuplexStream};
use tokio::net::{TcpListener, TcpStream};
use tokio::select;
use tokio::time::timeout;
#[derive(Debug, Clone)]
pub enum RaaSCmd {
Roll(u32),
Img(Vec<u8>),
}
use std::sync::Arc;
use tokio::io::{AsyncWriteExt, DuplexStream};
use tokio::sync::Mutex;
#[derive(Debug)]
struct ConnectorHandle {
pub id: u32,
pub stream: DuplexStream,
pub tags: Vec<raas::register::BotTypes>,
}
#[derive(Debug)]
struct ConnectorManager {
next_id: u32,
connectors: HashMap<u32, ConnectorHandle>,
}
#[derive(Debug, Clone)]
pub struct RobotManager {
addr: String,
connector_manager: std::sync::Arc<tokio::sync::Mutex<ConnectorManager>>,
command_stream: std::sync::Arc<tokio::sync::Mutex<DuplexStream>>,
command_stream: DuplexStream,
}
impl RobotManager {
pub fn new(addr: &str, stream: DuplexStream) -> Self {
let connector_manager = ConnectorManager {
next_id: 0,
connectors: Default::default(),
};
pub fn new(stream: DuplexStream) -> Self {
Self {
addr: addr.to_string(),
connector_manager: std::sync::Arc::new(tokio::sync::Mutex::new(connector_manager)),
command_stream: std::sync::Arc::new(tokio::sync::Mutex::new(stream)),
command_stream: stream,
}
}
async fn handle_register(&self, tcp_listener: &mut TcpListener) -> Result<(), std::io::Error> {
info!("Listening for connection at {}", self.addr);
let (mut socket, addr) = tcp_listener.accept().await?;
info!("Got connection from {}", addr);
let register_msg = recv_raas_msg(&mut socket).await?;
let register = Register::decode(&*register_msg.msg).unwrap();
let mut connector_manager = self.connector_manager.lock().await;
let id = connector_manager.next_id;
connector_manager.next_id += 1;
let register_resp = RegisterResponse {
name: register.name.to_string(),
r#type: register.bot_type,
id,
};
let mut msg = Vec::new();
register_resp.encode(&mut msg).unwrap();
send_raas_msg(&mut socket, msg).await?;
let (connector_duplex_manager, connector_duplex_connector) =
duplex(ROBOT_MESSAGE_QUEUE_SIZE);
let mut connector = RobotConnector::new(
id,
vec![register.bot_type()],
socket,
connector_duplex_connector,
);
let connector_handler = ConnectorHandle {
id,
stream: connector_duplex_manager,
tags: connector.bot_tags.clone(),
};
connector_manager.connectors.insert(id, connector_handler);
tokio::spawn(async move { connector.worker().await });
Ok(())
}
async fn recv_command(&self) -> Result<RaasMessage, Error> {
let mut command_stream = self.command_stream.lock().await;
let resp = recv_raas_msg(command_stream.deref_mut()).await?;
async fn recv_command(&mut self) -> Result<RaasMessage, Error> {
let resp = recv_raas_msg(&mut self.command_stream).await?;
Ok(resp)
}
async fn send_resp(&self, msg: RaasMessage) -> Result<(), Error> {
let mut command_stream = self.command_stream.lock().await;
command_stream.write_all(&*msg.into_bytes()).await?;
async fn send_resp(&mut self, msg: RaasMessage) -> Result<(), Error> {
self.command_stream.write_all(&msg.into_bytes()).await?;
Ok(())
}
pub async fn handle_command(&self) -> Result<(), Error> {
async fn handle_command(
&mut self,
connector_manager: Arc<Mutex<ConnectorManager>>,
) -> Result<(), Error> {
info!("Waiting for command...");
let msg = self.recv_command().await?;
@ -134,7 +47,7 @@ impl RobotManager {
Cmd::RollCmd(_) => raas::register::BotTypes::Roll,
};
let mut connector_manager = self.connector_manager.lock().await;
let mut connector_manager = connector_manager.lock().await;
let connector = connector_manager
.connectors
.iter_mut()
@ -144,7 +57,15 @@ impl RobotManager {
info!("Sending message to connector with id={}", id);
let mut msg = Vec::new();
request.encode(&mut msg)?;
send_raas_msg(&mut connector.stream, msg).await?;
let ret = send_raas_msg(&mut connector.stream, msg).await;
if let Err(err) = ret {
error!(
"Got '{}' sending message to Robot Connector id={}, closing connection",
err, id
);
return Err(Error::ConnectionClosed(*id));
}
let resp = recv_raas_msg(&mut connector.stream).await?;
info!("Got response from id={}", id);
@ -156,21 +77,22 @@ impl RobotManager {
Ok(())
}
pub async fn start_manager(self) -> Result<(), std::io::Error> {
let handle_command_manager = self.clone();
tokio::spawn(async move {
loop {
handle_command_manager.handle_command().await.unwrap()
}
});
pub async fn worker(
mut self,
connector_manager: Arc<Mutex<ConnectorManager>>,
) -> Result<(), Error> {
loop {
let res = self.handle_command(connector_manager.clone()).await;
tokio::spawn(async move {
let mut tcp_listener = TcpListener::bind(&self.addr).await.unwrap();
loop {
self.handle_register(&mut tcp_listener).await.unwrap()
}
});
if let Err(err) = res {
error!("Got error in robot manager: {}", err);
Ok(())
if let Error::ConnectionClosed(id) = err {
let mut connection_manager = connector_manager.lock().await;
connection_manager.connectors.remove(&id);
}
}
}
}
}