mirror of
https://github.com/Astatin3/unshell.git
synced 2026-06-08 22:38:01 -06:00
Make streams system
This commit is contained in:
+13
-1
@@ -40,7 +40,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
config: HashMap::from([(symbol!("host").to_string(), obs!("localhost:1234"))]),
|
config: HashMap::from([(symbol!("host").to_string(), obs!("localhost:1234"))]),
|
||||||
});
|
});
|
||||||
|
|
||||||
Manager::start_runtime(manager.clone(), runtime);
|
Manager::add_runtime(manager.clone(), runtime)?;
|
||||||
|
|
||||||
// Manager::st
|
// Manager::st
|
||||||
|
|
||||||
@@ -56,9 +56,21 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
|
|
||||||
match args[0] {
|
match args[0] {
|
||||||
"" => {}
|
"" => {}
|
||||||
|
"c" => {
|
||||||
|
println!(
|
||||||
|
"Current connections: {}",
|
||||||
|
manager.lock().unwrap().connections.len()
|
||||||
|
)
|
||||||
|
}
|
||||||
"test" => {
|
"test" => {
|
||||||
if let Some(arg) = args.get(1) {
|
if let Some(arg) = args.get(1) {
|
||||||
println!("Test with argument: {}", arg);
|
println!("Test with argument: {}", arg);
|
||||||
|
|
||||||
|
manager
|
||||||
|
.lock()
|
||||||
|
.unwrap()
|
||||||
|
.broadcast(unshell_lib::Announcement::TestAnnouncement(arg.to_string()))?;
|
||||||
|
|
||||||
// serverruntime
|
// serverruntime
|
||||||
// .send(&Announcement::TestAnnouncement(arg.to_string()))
|
// .send(&Announcement::TestAnnouncement(arg.to_string()))
|
||||||
// .unwrap();
|
// .unwrap();
|
||||||
|
|||||||
@@ -3,7 +3,7 @@ use bincode::{Decode, Encode};
|
|||||||
use crate::config::RuntimeConfig;
|
use crate::config::RuntimeConfig;
|
||||||
|
|
||||||
/// Mostly temporary server message type
|
/// Mostly temporary server message type
|
||||||
#[derive(Debug, Encode, Decode)]
|
#[derive(Clone, Debug, Encode, Decode)]
|
||||||
pub enum Announcement {
|
pub enum Announcement {
|
||||||
TestAnnouncement(String),
|
TestAnnouncement(String),
|
||||||
|
|
||||||
|
|||||||
@@ -1,5 +1,4 @@
|
|||||||
use std::{
|
use std::{
|
||||||
io::Read,
|
|
||||||
net::TcpStream,
|
net::TcpStream,
|
||||||
sync::{
|
sync::{
|
||||||
Arc,
|
Arc,
|
||||||
@@ -9,83 +8,24 @@ use std::{
|
|||||||
time::Duration,
|
time::Duration,
|
||||||
};
|
};
|
||||||
|
|
||||||
use crate::{config::RuntimeConfig, *};
|
use crate::{config::RuntimeConfig, network::Stream, *};
|
||||||
// use unshell_modules::{Manager, ModuleRuntime};
|
// use unshell_modules::{Manager, ModuleRuntime};
|
||||||
|
|
||||||
use crate::{Announcement, ModuleRuntime};
|
use crate::ModuleRuntime;
|
||||||
|
|
||||||
pub struct ClientRuntime {
|
pub struct ClientRuntime {
|
||||||
thread_handle: JoinHandle<()>,
|
config: &'static RuntimeConfig,
|
||||||
|
thread_handle: Option<JoinHandle<()>>,
|
||||||
join_signal: Arc<AtomicBool>,
|
join_signal: Arc<AtomicBool>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ClientRuntime {
|
impl ClientRuntime {
|
||||||
pub fn new(config: &'static RuntimeConfig) -> Result<ClientRuntime, ModuleError> {
|
pub fn new(config: &'static RuntimeConfig) -> Result<ClientRuntime, ModuleError> {
|
||||||
let join_signal = Arc::new(AtomicBool::new(false));
|
let join_signal = Arc::new(AtomicBool::new(false));
|
||||||
let join_clone = join_signal.clone();
|
|
||||||
|
|
||||||
let host = match config.config.get("host") {
|
|
||||||
Some(host) => host,
|
|
||||||
None => {
|
|
||||||
return Err(ModuleError::Error(
|
|
||||||
"Could not find HOST in Client Runtime".into(),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let retry = match config.config.get("retry") {
|
|
||||||
Some(host) => Duration::from_millis(host.parse::<u64>().unwrap()),
|
|
||||||
None => {
|
|
||||||
return Err(ModuleError::Error(
|
|
||||||
"Could not find RETRY in Client Runtime".into(),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
Ok(Self {
|
Ok(Self {
|
||||||
thread_handle: thread::spawn(move || {
|
config,
|
||||||
debug!("Connecting to server...");
|
thread_handle: None,
|
||||||
|
|
||||||
loop {
|
|
||||||
let mut stream = match TcpStream::connect(host) {
|
|
||||||
Ok(stream) => stream,
|
|
||||||
Err(e) => {
|
|
||||||
error!("Failed to connect to server: {}", e);
|
|
||||||
thread::sleep(retry);
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
info!("Connected");
|
|
||||||
|
|
||||||
while !join_clone.load(Ordering::Relaxed) {
|
|
||||||
let mut size_buf = [0u8; 4];
|
|
||||||
match stream.read_exact(&mut size_buf) {
|
|
||||||
Ok(()) => {}
|
|
||||||
Err(_) => {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
let size = u32::from_be_bytes(size_buf);
|
|
||||||
|
|
||||||
let mut buf = vec![0u8; size as usize];
|
|
||||||
|
|
||||||
stream.read_exact(&mut buf).unwrap();
|
|
||||||
|
|
||||||
let a = Announcement::decode(&buf).unwrap();
|
|
||||||
|
|
||||||
match a {
|
|
||||||
Announcement::TestAnnouncement(s) => {
|
|
||||||
println!("Received test announcement: {}", s)
|
|
||||||
}
|
|
||||||
_ => {}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
debug!("Disconnected from {}", host);
|
|
||||||
|
|
||||||
thread::sleep(retry);
|
|
||||||
}
|
|
||||||
}),
|
|
||||||
join_signal,
|
join_signal,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
@@ -109,13 +49,75 @@ impl ClientRuntime {
|
|||||||
|
|
||||||
impl ModuleRuntime for ClientRuntime {
|
impl ModuleRuntime for ClientRuntime {
|
||||||
fn is_running(&self) -> bool {
|
fn is_running(&self) -> bool {
|
||||||
!self.thread_handle.is_finished()
|
self.thread_handle.as_ref().is_none_or(|h| h.is_finished())
|
||||||
}
|
}
|
||||||
|
|
||||||
fn kill(self: Box<Self>) {
|
fn kill(self: Box<Self>) {
|
||||||
if !self.thread_handle.is_finished() {
|
if !self.is_running() {
|
||||||
self.join_signal.store(true, Ordering::Relaxed);
|
self.join_signal.store(true, Ordering::Relaxed);
|
||||||
let _ = self.thread_handle.join();
|
if let Some(handle) = self.thread_handle {
|
||||||
|
let _ = handle.join();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn init(&mut self, manager: Arc<Mutex<Manager>>) -> Result<(), ModuleError> {
|
||||||
|
let host = match self.config.config.get("host") {
|
||||||
|
Some(host) => host,
|
||||||
|
None => {
|
||||||
|
return Err(ModuleError::Error(
|
||||||
|
"Could not find HOST in Client Runtime".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let retry = match self.config.config.get("retry") {
|
||||||
|
Some(retry) => Duration::from_millis(retry.parse::<u64>().unwrap()),
|
||||||
|
None => {
|
||||||
|
return Err(ModuleError::Error(
|
||||||
|
"Could not find RETRY in Client Runtime".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// let join_clone = self.join_signal.clone();
|
||||||
|
|
||||||
|
thread::spawn(move || {
|
||||||
|
debug!("Connecting to server...");
|
||||||
|
|
||||||
|
loop {
|
||||||
|
let stream = match TcpStream::connect(host) {
|
||||||
|
Ok(stream) => stream,
|
||||||
|
Err(e) => {
|
||||||
|
error!("Failed to connect to server: {}", e);
|
||||||
|
thread::sleep(retry);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
info!("Connected to {}", host);
|
||||||
|
|
||||||
|
thread::sleep(Duration::from_millis(100));
|
||||||
|
// Duration::from_millis(100);
|
||||||
|
|
||||||
|
let stream = crate::network::TcpStream::new(stream);
|
||||||
|
let stream_clone = stream.try_clone().unwrap();
|
||||||
|
|
||||||
|
manager.lock().unwrap().add_connection(stream_clone);
|
||||||
|
|
||||||
|
// while !join_clone.load(Ordering::Relaxed) {
|
||||||
|
|
||||||
|
// }
|
||||||
|
|
||||||
|
while stream.is_alive() {
|
||||||
|
thread::sleep(Duration::from_millis(100));
|
||||||
|
}
|
||||||
|
|
||||||
|
debug!("Disconnected from 1234 {}", host);
|
||||||
|
|
||||||
|
thread::sleep(retry);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -11,10 +11,15 @@ mod components;
|
|||||||
pub use components::get_components;
|
pub use components::get_components;
|
||||||
|
|
||||||
mod announcement;
|
mod announcement;
|
||||||
use std::fmt::{self, Debug};
|
use std::{
|
||||||
|
fmt::{self, Debug},
|
||||||
|
sync::{Arc, Mutex},
|
||||||
|
};
|
||||||
|
|
||||||
pub use announcement::Announcement;
|
pub use announcement::Announcement;
|
||||||
|
|
||||||
|
use crate::module::Manager;
|
||||||
|
|
||||||
///Generic error type for module-related operations.
|
///Generic error type for module-related operations.
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
pub enum ModuleError {
|
pub enum ModuleError {
|
||||||
@@ -47,6 +52,8 @@ impl fmt::Display for ModuleError {
|
|||||||
|
|
||||||
/// Trait for defining modules that have a runtime.
|
/// Trait for defining modules that have a runtime.
|
||||||
pub trait ModuleRuntime: Send + Sync {
|
pub trait ModuleRuntime: Send + Sync {
|
||||||
|
fn init(&mut self, manager: Arc<Mutex<Manager>>) -> Result<(), ModuleError>;
|
||||||
|
|
||||||
/// Returns true if the module is running.
|
/// Returns true if the module is running.
|
||||||
/// After returning false, the module will be dropped.
|
/// After returning false, the module will be dropped.
|
||||||
fn is_running(&self) -> bool;
|
fn is_running(&self) -> bool;
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ use std::{
|
|||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
config::{NamedComponent, PayloadConfig, RuntimeConfig},
|
config::{NamedComponent, PayloadConfig, RuntimeConfig},
|
||||||
network::Connection,
|
network::Stream,
|
||||||
*,
|
*,
|
||||||
};
|
};
|
||||||
use module::Module;
|
use module::Module;
|
||||||
@@ -24,7 +24,7 @@ pub struct Manager {
|
|||||||
components: HashMap<String, NamedComponent>,
|
components: HashMap<String, NamedComponent>,
|
||||||
active_runtimes: Vec<Box<dyn ModuleRuntime>>,
|
active_runtimes: Vec<Box<dyn ModuleRuntime>>,
|
||||||
|
|
||||||
pub connections: Vec<Connection>,
|
pub connections: Vec<Box<dyn Stream<Announcement>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
// static mut MANAGER_RUNTIME: Option<Arc<Mutex<Manager>>> = None;
|
// static mut MANAGER_RUNTIME: Option<Arc<Mutex<Manager>>> = None;
|
||||||
@@ -60,9 +60,16 @@ impl Manager {
|
|||||||
|
|
||||||
let this = Arc::new(Mutex::new(this));
|
let this = Arc::new(Mutex::new(this));
|
||||||
|
|
||||||
debug!("Starting runtimes...");
|
debug!("Creating runtimes...");
|
||||||
for runtime in &config.runtime_config {
|
for runtime in &config.runtime_config {
|
||||||
Self::start_runtime(this.clone(), runtime);
|
Self::create_runtime(this.clone(), runtime);
|
||||||
|
}
|
||||||
|
|
||||||
|
debug!("Starting runtimes...");
|
||||||
|
for runtime in &mut this.lock().unwrap().active_runtimes {
|
||||||
|
if let Err(e) = runtime.init(this.clone()) {
|
||||||
|
warn!("Failed to start runtime: {}", e);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
this.lock().unwrap().handle = Some(Self::start_thread(this.clone()));
|
this.lock().unwrap().handle = Some(Self::start_thread(this.clone()));
|
||||||
@@ -141,7 +148,7 @@ impl Manager {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Start a runtime
|
/// Start a runtime
|
||||||
pub fn start_runtime<'a>(this: Arc<Mutex<Self>>, runtime: &'static RuntimeConfig) {
|
fn create_runtime<'a>(this: Arc<Mutex<Self>>, runtime: &'static RuntimeConfig) {
|
||||||
let mut this_lock = this.lock().unwrap();
|
let mut this_lock = this.lock().unwrap();
|
||||||
|
|
||||||
let component = match this_lock.components.get(&runtime.parent_component) {
|
let component = match this_lock.components.get(&runtime.parent_component) {
|
||||||
@@ -168,6 +175,21 @@ impl Manager {
|
|||||||
this_lock.active_runtimes.push(runtime);
|
this_lock.active_runtimes.push(runtime);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn add_runtime(
|
||||||
|
this: Arc<Mutex<Self>>,
|
||||||
|
runtime: &'static RuntimeConfig,
|
||||||
|
) -> Result<(), ModuleError> {
|
||||||
|
Self::create_runtime(this.clone(), runtime);
|
||||||
|
|
||||||
|
this.lock()
|
||||||
|
.unwrap()
|
||||||
|
.active_runtimes
|
||||||
|
.iter_mut()
|
||||||
|
.last()
|
||||||
|
.unwrap()
|
||||||
|
.init(this.clone())
|
||||||
|
}
|
||||||
|
|
||||||
pub fn get_name(&self) -> &str {
|
pub fn get_name(&self) -> &str {
|
||||||
self.id
|
self.id
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,11 +1,7 @@
|
|||||||
use crate::{
|
use crate::{Announcement, ModuleError, module::Manager, network::Stream};
|
||||||
Announcement,
|
|
||||||
module::Manager,
|
|
||||||
network::{Connection, Stream},
|
|
||||||
};
|
|
||||||
|
|
||||||
impl Manager {
|
impl Manager {
|
||||||
pub fn add_connection(&mut self, connection: Connection) {
|
pub fn add_connection(&mut self, connection: Box<dyn Stream<Announcement>>) {
|
||||||
self.connections.push(connection);
|
self.connections.push(connection);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -17,8 +13,8 @@ impl Manager {
|
|||||||
// Collect all incoming announcements
|
// Collect all incoming announcements
|
||||||
let announcements = self
|
let announcements = self
|
||||||
.connections
|
.connections
|
||||||
.iter()
|
.iter_mut()
|
||||||
.map(|c| c.read())
|
.map(|c| c.try_read())
|
||||||
.flat_map(|array| array)
|
.flat_map(|array| array)
|
||||||
.collect::<Vec<Announcement>>();
|
.collect::<Vec<Announcement>>();
|
||||||
|
|
||||||
@@ -26,4 +22,11 @@ impl Manager {
|
|||||||
self.recv_announcement(&announcement)
|
self.recv_announcement(&announcement)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn broadcast(&mut self, announcement: Announcement) -> Result<(), ModuleError> {
|
||||||
|
for connection in &mut self.connections {
|
||||||
|
connection.write(announcement.clone())?;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,65 +1,65 @@
|
|||||||
use std::sync::{
|
// use std::sync::{
|
||||||
Arc,
|
// Arc,
|
||||||
atomic::{AtomicBool, Ordering},
|
// atomic::{AtomicBool, Ordering},
|
||||||
};
|
// };
|
||||||
|
|
||||||
use crate::{Announcement, ModuleError, network::Stream};
|
// use crate::{Announcement, ModuleError, network::Stream};
|
||||||
|
|
||||||
use crossbeam_channel::{Receiver, Sender};
|
// use crossbeam_channel::{Receiver, Sender};
|
||||||
|
|
||||||
pub struct Connection {
|
// pub struct Connection {
|
||||||
tx: Sender<Announcement>,
|
// tx: Sender<Announcement>,
|
||||||
rx: Receiver<Announcement>,
|
// rx: Receiver<Announcement>,
|
||||||
is_alive: Arc<AtomicBool>,
|
// is_alive: Arc<AtomicBool>,
|
||||||
}
|
// }
|
||||||
|
|
||||||
impl Connection {
|
// impl Connection {
|
||||||
pub fn new() -> (Connection, Connection) {
|
// pub fn new() -> (Connection, Connection) {
|
||||||
let (tx_mgr, rx) = crossbeam_channel::unbounded();
|
// let (tx_mgr, rx) = crossbeam_channel::unbounded();
|
||||||
let (tx, rx_mgr) = crossbeam_channel::unbounded();
|
// let (tx, rx_mgr) = crossbeam_channel::unbounded();
|
||||||
let alive = Arc::new(AtomicBool::new(false));
|
// let alive = Arc::new(AtomicBool::new(false));
|
||||||
|
|
||||||
(
|
// (
|
||||||
Self {
|
// Self {
|
||||||
tx: tx_mgr,
|
// tx: tx_mgr,
|
||||||
rx: rx_mgr,
|
// rx: rx_mgr,
|
||||||
is_alive: alive.clone(),
|
// is_alive: alive.clone(),
|
||||||
},
|
// },
|
||||||
Self {
|
// Self {
|
||||||
tx,
|
// tx,
|
||||||
rx,
|
// rx,
|
||||||
is_alive: alive,
|
// is_alive: alive,
|
||||||
},
|
// },
|
||||||
)
|
// )
|
||||||
}
|
// }
|
||||||
}
|
// }
|
||||||
|
|
||||||
impl Stream<Announcement> for Connection {
|
// impl Stream<Announcement> for Connection {
|
||||||
fn is_alive(&self) -> bool {
|
// fn is_alive(&self) -> bool {
|
||||||
self.is_alive.load(Ordering::Relaxed)
|
// self.is_alive.load(Ordering::Relaxed)
|
||||||
}
|
// }
|
||||||
|
|
||||||
fn len(&self) -> usize {
|
// fn len(&self) -> usize {
|
||||||
self.rx.len()
|
// self.rx.len()
|
||||||
}
|
// }
|
||||||
|
|
||||||
fn read(&self) -> Vec<Announcement> {
|
// fn read(&self) -> Vec<Announcement> {
|
||||||
self.rx.try_iter().collect()
|
// self.rx.try_iter().collect()
|
||||||
}
|
// }
|
||||||
|
|
||||||
fn write(&mut self, data: Announcement) -> Result<(), crate::ModuleError> {
|
// fn write(&mut self, data: Announcement) -> Result<(), crate::ModuleError> {
|
||||||
self.tx
|
// self.tx
|
||||||
.send(data)
|
// .send(data)
|
||||||
.map_err(|_| ModuleError::Error("Failed to send".into()))?;
|
// .map_err(|_| ModuleError::Error("Failed to send".into()))?;
|
||||||
|
|
||||||
Ok(())
|
// Ok(())
|
||||||
}
|
// }
|
||||||
|
|
||||||
fn try_clone(&self) -> Result<Box<dyn Stream<Announcement> + Send + Sync>, crate::ModuleError> {
|
// fn try_clone(&self) -> Result<Box<dyn Stream<Announcement> + Send + Sync>, crate::ModuleError> {
|
||||||
Ok(Box::new(Self {
|
// Ok(Box::new(Self {
|
||||||
tx: self.tx.clone(),
|
// tx: self.tx.clone(),
|
||||||
rx: self.rx.clone(),
|
// rx: self.rx.clone(),
|
||||||
is_alive: self.is_alive.clone(),
|
// is_alive: self.is_alive.clone(),
|
||||||
}))
|
// }))
|
||||||
}
|
// }
|
||||||
}
|
// }
|
||||||
|
|||||||
@@ -1,6 +1,8 @@
|
|||||||
mod connection;
|
// mod connection;
|
||||||
|
mod tcp_stream;
|
||||||
|
pub use tcp_stream::TcpStream;
|
||||||
|
|
||||||
pub use connection::Connection;
|
// pub use connection::Connection;
|
||||||
|
|
||||||
use crate::ModuleError;
|
use crate::ModuleError;
|
||||||
|
|
||||||
@@ -9,8 +11,19 @@ pub trait Stream<T>: Send + Sync {
|
|||||||
// fn get_info(&self) -> String;
|
// fn get_info(&self) -> String;
|
||||||
fn is_alive(&self) -> bool;
|
fn is_alive(&self) -> bool;
|
||||||
|
|
||||||
fn len(&self) -> usize;
|
fn has_recv(&self) -> bool;
|
||||||
fn read(&self) -> Vec<T>;
|
|
||||||
|
/// Possibly blocking stream read function
|
||||||
|
fn read(&mut self) -> Vec<T>;
|
||||||
|
|
||||||
|
/// Non-blocking read function
|
||||||
|
fn try_read(&mut self) -> Vec<T> {
|
||||||
|
if self.has_recv() {
|
||||||
|
self.read()
|
||||||
|
} else {
|
||||||
|
Vec::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn write(&mut self, data: T) -> Result<(), ModuleError>;
|
fn write(&mut self, data: T) -> Result<(), ModuleError>;
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,126 @@
|
|||||||
|
use std::{
|
||||||
|
io::{Read, Write},
|
||||||
|
net,
|
||||||
|
sync::{
|
||||||
|
Arc,
|
||||||
|
atomic::{AtomicBool, Ordering},
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
use crate::{Announcement, ModuleError, debug, network::Stream};
|
||||||
|
|
||||||
|
pub struct TcpStream(Arc<AtomicBool>, net::TcpStream);
|
||||||
|
|
||||||
|
impl TcpStream {
|
||||||
|
pub fn new(stream: net::TcpStream) -> Self {
|
||||||
|
stream.set_nonblocking(true).unwrap();
|
||||||
|
Self(Arc::new(AtomicBool::new(true)), stream)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Call this when the stream ends
|
||||||
|
fn disconnected(&mut self) {
|
||||||
|
self.0.store(false, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Stream<Announcement> for TcpStream {
|
||||||
|
fn is_alive(&self) -> bool {
|
||||||
|
// if self.1.take_error().unwrap_or(None).is_some() {
|
||||||
|
// // self.1.pe
|
||||||
|
// warn!("Disconnected #################");
|
||||||
|
// return true;
|
||||||
|
// } else {
|
||||||
|
// return false;
|
||||||
|
// }
|
||||||
|
|
||||||
|
// let mut buf = [0u8; 1];
|
||||||
|
// match self.1.peek(&mut buf) {
|
||||||
|
// Ok(n) => n == 1,
|
||||||
|
// Err(_) => false,
|
||||||
|
// }
|
||||||
|
|
||||||
|
let mut buf = [0u8; 1];
|
||||||
|
match self.1.peek(&mut buf) {
|
||||||
|
Ok(0) => false, // Connection closed (EOF)
|
||||||
|
Ok(_) => true, // Data available or connection alive
|
||||||
|
Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => true, // No data but alive
|
||||||
|
Err(_) => false, // Connection error
|
||||||
|
}
|
||||||
|
|
||||||
|
// true
|
||||||
|
|
||||||
|
// self.0.load(Ordering::Relaxed)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn has_recv(&self) -> bool {
|
||||||
|
let mut buf = [0u8; 1];
|
||||||
|
match self.1.peek(&mut buf) {
|
||||||
|
Ok(n) if n > 0 => true, // Data is available
|
||||||
|
Ok(_) => false, // EOF (connection closed)
|
||||||
|
Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => false, // No data
|
||||||
|
Err(_) => false,
|
||||||
|
}
|
||||||
|
// false
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read(&mut self) -> Vec<Announcement> {
|
||||||
|
let mut ret = Vec::new();
|
||||||
|
|
||||||
|
while self.has_recv() {
|
||||||
|
let mut size_buf = [0u8; 4];
|
||||||
|
match self.1.read_exact(&mut size_buf) {
|
||||||
|
Ok(()) => {}
|
||||||
|
Err(_) => {
|
||||||
|
self.disconnected();
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let size = u32::from_be_bytes(size_buf);
|
||||||
|
|
||||||
|
let mut buf = vec![0u8; size as usize];
|
||||||
|
|
||||||
|
match self.1.read_exact(&mut buf) {
|
||||||
|
Ok(()) => {}
|
||||||
|
Err(_) => {
|
||||||
|
self.disconnected();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(a) = Announcement::decode(&buf) {
|
||||||
|
ret.push(a);
|
||||||
|
} else {
|
||||||
|
debug!("Malformed data");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
ret
|
||||||
|
}
|
||||||
|
|
||||||
|
fn write(&mut self, announcement: Announcement) -> Result<(), crate::ModuleError> {
|
||||||
|
let bytes = announcement.encode();
|
||||||
|
|
||||||
|
// Write length of bytes
|
||||||
|
self.1
|
||||||
|
.write_all(&u32::to_be_bytes(bytes.len() as u32))
|
||||||
|
.map_err(|e| ModuleError::Error(e.to_string().into()))?;
|
||||||
|
// Write data
|
||||||
|
self.1
|
||||||
|
.write_all(&bytes)
|
||||||
|
.map_err(|e| ModuleError::Error(e.to_string().into()))?;
|
||||||
|
// Flush data
|
||||||
|
self.1
|
||||||
|
.flush()
|
||||||
|
.map_err(|e| ModuleError::Error(e.to_string().into()))?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn try_clone(&self) -> Result<Box<dyn Stream<Announcement> + Send + Sync>, crate::ModuleError> {
|
||||||
|
Ok(Box::new(Self(
|
||||||
|
self.0.clone(),
|
||||||
|
self.1
|
||||||
|
.try_clone()
|
||||||
|
.map_err(|e| ModuleError::Error(e.to_string().into()))?,
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,26 +1,66 @@
|
|||||||
use std::{
|
use std::{
|
||||||
io::{Read, Write},
|
net::TcpListener,
|
||||||
net::{TcpListener, TcpStream},
|
|
||||||
sync::{Arc, Mutex},
|
sync::{Arc, Mutex},
|
||||||
thread::{self, JoinHandle},
|
thread::{self, JoinHandle},
|
||||||
};
|
};
|
||||||
|
|
||||||
use crate::{config::RuntimeConfig, *};
|
use crate::{config::RuntimeConfig, module::Manager, *};
|
||||||
|
|
||||||
pub struct ListenerRuntime {
|
pub struct ListenerRuntime {
|
||||||
thread_handle: JoinHandle<()>,
|
config: &'static RuntimeConfig,
|
||||||
// join_signal: Arc<AtomicBool>,
|
thread_handle: Option<JoinHandle<()>>,
|
||||||
// listener: TcpListener,
|
// streams: Arc<Mutex<Vec<TcpStream>>>,
|
||||||
streams: Arc<Mutex<Vec<TcpStream>>>,
|
// manager: Option<Arc<Mutex<Manager>>>,
|
||||||
// reader: BufReader<TcpListener>,
|
|
||||||
// writer: BufWriter<TcpListener>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ListenerRuntime {
|
impl ListenerRuntime {
|
||||||
pub fn new(config: &'static RuntimeConfig) -> Result<Self, ModuleError> {
|
pub fn new(config: &'static RuntimeConfig) -> Result<Self, ModuleError> {
|
||||||
// info!("Starting listener runtime on {}",);
|
Ok(Self {
|
||||||
|
config,
|
||||||
|
thread_handle: None,
|
||||||
|
// streams: Arc::new(Mutex::new(Vec::new())),
|
||||||
|
// manager: None,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
let host = match config.config.get("host") {
|
// pub fn send(&mut self, announcement: &Announcement) -> Result<(), std::io::Error> {
|
||||||
|
// let bytes = announcement.encode();
|
||||||
|
|
||||||
|
// let mut streams = self.streams.lock().unwrap();
|
||||||
|
|
||||||
|
// for stream in streams.iter_mut() {
|
||||||
|
// stream.write_all(&u32::to_be_bytes(bytes.len() as u32))?;
|
||||||
|
// stream.write_all(&bytes)?;
|
||||||
|
// stream.flush()?;
|
||||||
|
// }
|
||||||
|
|
||||||
|
// debug!("Announcement {:?} sent", announcement);
|
||||||
|
|
||||||
|
// Ok(())
|
||||||
|
// }
|
||||||
|
|
||||||
|
// pub fn recv(&mut self) -> Result<Announcement, ModuleError> {
|
||||||
|
// let stream = &mut self.streams.lock().unwrap()[0];
|
||||||
|
|
||||||
|
// let mut size_buf = [0u8; 4];
|
||||||
|
// stream.read_exact(&mut size_buf).unwrap();
|
||||||
|
// let size = u32::from_be_bytes(size_buf);
|
||||||
|
|
||||||
|
// let mut buf = vec![0u8; size as usize];
|
||||||
|
|
||||||
|
// stream.read_exact(&mut buf).unwrap();
|
||||||
|
|
||||||
|
// if let Some(announcement) = Announcement::decode(&buf) {
|
||||||
|
// Ok(announcement)
|
||||||
|
// } else {
|
||||||
|
// Err(ModuleError::Error("Failed to decode announcement".into()))
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ModuleRuntime for ListenerRuntime {
|
||||||
|
fn init(&mut self, manager: Arc<Mutex<Manager>>) -> Result<(), ModuleError> {
|
||||||
|
let host = match self.config.config.get("host") {
|
||||||
Some(host) => host,
|
Some(host) => host,
|
||||||
None => {
|
None => {
|
||||||
return Err(ModuleError::Error(
|
return Err(ModuleError::Error(
|
||||||
@@ -30,71 +70,39 @@ impl ListenerRuntime {
|
|||||||
};
|
};
|
||||||
|
|
||||||
let listener = TcpListener::bind(host).unwrap();
|
let listener = TcpListener::bind(host).unwrap();
|
||||||
let streams = Arc::new(Mutex::new(Vec::new()));
|
// let streams = Arc::new(Mutex::new(Vec::new()));
|
||||||
|
|
||||||
let streams_clone = streams.clone();
|
// let streams_clone = streams.clone();
|
||||||
|
|
||||||
let thread_handle = thread::spawn(move || {
|
let thread_handle = thread::spawn(move || {
|
||||||
let streams = streams_clone.clone();
|
// let streams = streams_clone.clone();
|
||||||
for stream in listener.incoming() {
|
for stream in listener.incoming() {
|
||||||
let stream = stream.unwrap();
|
let stream = stream.unwrap();
|
||||||
debug!("New connection from {}", stream.peer_addr().unwrap());
|
debug!("New connection from {}", stream.peer_addr().unwrap());
|
||||||
streams.lock().unwrap().push(stream);
|
|
||||||
|
let stream = crate::network::TcpStream::new(stream);
|
||||||
|
|
||||||
|
manager.lock().unwrap().add_connection(Box::new(stream));
|
||||||
|
|
||||||
|
// streams.lock().unwrap().push(stream);
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
Ok(Self {
|
|
||||||
thread_handle,
|
|
||||||
streams,
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn send(&mut self, announcement: &Announcement) -> Result<(), std::io::Error> {
|
self.thread_handle = Some(thread_handle);
|
||||||
let bytes = announcement.encode();
|
|
||||||
|
|
||||||
let mut streams = self.streams.lock().unwrap();
|
|
||||||
|
|
||||||
for stream in streams.iter_mut() {
|
|
||||||
stream.write_all(&u32::to_be_bytes(bytes.len() as u32))?;
|
|
||||||
stream.write_all(&bytes)?;
|
|
||||||
stream.flush()?;
|
|
||||||
}
|
|
||||||
|
|
||||||
debug!("Announcement {:?} sent", announcement);
|
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn recv(&mut self) -> Result<Announcement, ModuleError> {
|
|
||||||
let stream = &mut self.streams.lock().unwrap()[0];
|
|
||||||
|
|
||||||
let mut size_buf = [0u8; 4];
|
|
||||||
stream.read_exact(&mut size_buf).unwrap();
|
|
||||||
let size = u32::from_be_bytes(size_buf);
|
|
||||||
|
|
||||||
let mut buf = vec![0u8; size as usize];
|
|
||||||
|
|
||||||
stream.read_exact(&mut buf).unwrap();
|
|
||||||
|
|
||||||
if let Some(announcement) = Announcement::decode(&buf) {
|
|
||||||
Ok(announcement)
|
|
||||||
} else {
|
|
||||||
Err(ModuleError::Error("Failed to decode announcement".into()))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl ModuleRuntime for ListenerRuntime {
|
|
||||||
// fn init(&mut self) {}
|
|
||||||
|
|
||||||
fn is_running(&self) -> bool {
|
fn is_running(&self) -> bool {
|
||||||
true
|
true
|
||||||
}
|
}
|
||||||
|
|
||||||
fn kill(self: Box<Self>) {
|
fn kill(self: Box<Self>) {
|
||||||
if !self.thread_handle.is_finished() {
|
// if let Some(thread)
|
||||||
// self.join_signal.store(true, Ordering::Relaxed);
|
// if !self.thread_handle.is_finished() {
|
||||||
let _ = self.thread_handle.join();
|
// // self.join_signal.store(true, Ordering::Relaxed);
|
||||||
}
|
// let _ = self.thread_handle.join();
|
||||||
// drop(self);
|
// }
|
||||||
|
// // drop(self);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user