working

Author: Autoparallel

arbiter-core/src/agent.rs

@@ -15,7 +15,15 @@ pub struct Agent<L: LifeCycle, C: Connection> {
   inner:                 L,
   pub(crate) controller: Arc<Controller>,
   handlers:              HashMap<TypeId, MessageHandlerFn<C>>,
-  transport:             Transport<C>,
+  pub(crate) transport:  Transport<C>,
+}
+
+pub struct RunningAgent<C: Connection> {
+  pub name: Option<String>,
+  pub(crate) task: JoinHandle<Box<dyn RuntimeAgent<C>>>,
+  pub(crate) controller: Arc<Controller>,
+  pub(crate) outbound_connections: Arc<Mutex<HashMap<C::Address, C::Sender>>>,
+  pub(crate) sender: C::Sender,
 }
 
 pub struct Controller {
@@ -86,11 +94,6 @@ impl<L: LifeCycle, C: Connection> Agent<L, C> {
 
   pub fn clear_name(&mut self) { self.name = None; }
 
-  // Access to inner L, requires locking.
-  pub fn inner_mut(&mut self) -> &mut L { &mut self.inner }
-
-  pub fn inner(&self) -> &L { &self.inner }
-
   pub fn with_handler<M>(mut self) -> Self
   where
     M: Message,
@@ -137,10 +140,8 @@ pub trait RuntimeAgent<C: Connection>: Send + Sync + Any {
   // Methods to signal the agent's desired state
   fn signal_start(&self);
   fn signal_stop(&self);
-  fn current_loop_state(&self) -> State;
 
-  fn process(self) -> JoinHandle<Self>
-  where Self: Sized;
+  fn process(self) -> RunningAgent<C>;
 
   #[cfg(test)]
   fn inner_as_any(&self) -> &dyn Any;
@@ -175,12 +176,16 @@ where C::Payload: Package<L::StartMessage> + Package<L::StopMessage>
     }
   }
 
-  fn current_loop_state(&self) -> State { *self.controller.state.lock().unwrap() }
+  fn process(mut self) -> RunningAgent<C> {
+    let name = self.name.clone();
+    let outbound_connections = self.transport.outbound_connections.clone();
+    let controller = self.controller.clone();
+    let (_, sender) = self.transport.create_inbound_connection();
 
-  fn process(mut self) -> JoinHandle<Self> {
-    std::thread::spawn(move || {
+    let handle = std::thread::spawn(move || {
       let start_message = self.inner.on_start();
       self.transport.broadcast(Envelope::package(start_message));
+      println!("Agent {}: Sent start message.", self.address());
 
       loop {
         println!("Agent {}: Loop iteration.", self.address());
@@ -269,8 +274,9 @@ where C::Payload: Package<L::StartMessage> + Package<L::StopMessage>
 
       self.inner.on_stop();
       println!("Agent {}: Executed on_stop. Main loop finished.", self.address());
-      self
-    })
+      Box::new(self) as Box<dyn RuntimeAgent<C>>
+    });
+    RunningAgent { name, task: handle, controller, outbound_connections, sender }
   }
 
   #[cfg(test)]
@@ -287,20 +293,21 @@ mod tests {
   fn test_agent_lifecycle() {
     let logger = Logger { name: "TestLogger".to_string(), message_count: 0 };
     let agent = Agent::<Logger, InMemory>::new(logger);
-    assert_eq!(agent.current_loop_state(), State::Stopped);
+    assert_eq!(*agent.controller.state.lock().unwrap(), State::Stopped);
 
     // Start agent
     let controller = agent.controller.clone();
-    let handle = agent.process();
+    let running_agent = agent.process();
 
     controller.signal_start();
     std::thread::sleep(std::time::Duration::from_millis(50)); // Give time for agent to start
     assert_eq!(controller.state.lock().unwrap().clone(), State::Running);
 
     controller.signal_stop();
-    let joined_agent = handle.join().unwrap();
-    assert_eq!(joined_agent.inner().message_count, 0);
-    assert_eq!(joined_agent.current_loop_state(), State::Stopped);
+    let joined_agent = running_agent.task.join().unwrap();
+    let agent = joined_agent.inner_as_any().downcast_ref::<Logger>().unwrap();
+    assert_eq!(agent.message_count, 0);
+    assert_eq!(*running_agent.controller.state.lock().unwrap(), State::Stopped);
   }
 
   #[test]
@@ -310,7 +317,7 @@ mod tests {
     let sender = agent.transport.inbound_connection.sender.clone();
 
     let controller = agent.controller.clone();
-    let handle = agent.process();
+    let running_agent = agent.process();
 
     controller.signal_start();
     std::thread::sleep(std::time::Duration::from_millis(50)); // Allow agent to start and enter wait
@@ -321,9 +328,9 @@ mod tests {
     assert_eq!(controller.state.lock().unwrap().clone(), State::Running);
 
     controller.signal_stop();
-    let result_agent = handle.join().unwrap();
-
-    assert_eq!(result_agent.inner().message_count, 1);
+    let result_agent = running_agent.task.join().unwrap();
+    let agent = result_agent.inner_as_any().downcast_ref::<Logger>().unwrap();
+    assert_eq!(agent.message_count, 1);
   }
 
   #[test]
@@ -335,10 +342,10 @@ mod tests {
     agent_struct = agent_struct.with_handler::<TextMessage>().with_handler::<NumberMessage>();
     let sender = agent_struct.transport.inbound_connection.sender.clone();
 
-    assert_eq!(agent_struct.current_loop_state(), State::Stopped);
+    assert_eq!(*agent_struct.controller.state.lock().unwrap(), State::Stopped);
 
     let controller = agent_struct.controller.clone();
-    let handle = agent_struct.process();
+    let running_agent = agent_struct.process();
 
     controller.signal_start();
     std::thread::sleep(std::time::Duration::from_millis(50));
@@ -352,7 +359,8 @@ mod tests {
     std::thread::sleep(std::time::Duration::from_millis(50));
 
     controller.signal_stop();
-    let result_agent = handle.join().unwrap();
-    assert_eq!(result_agent.inner().message_count, 2);
+    let result_agent = running_agent.task.join().unwrap();
+    let agent = result_agent.inner_as_any().downcast_ref::<Logger>().unwrap();
+    assert_eq!(agent.message_count, 2);
   }
 }

arbiter-core/src/connection/memory.rs

@@ -2,7 +2,7 @@ use std::sync::Arc;
 
 use crate::{
   agent::AgentIdentity,
-  connection::{Connection, Receiver, Sender},
+  connection::{Connection, GetNew, Receiver, Sender},
   handler::{Envelope, Message},
 };
 
@@ -13,6 +13,10 @@ pub struct InMemory {
   receiver:           flume::Receiver<Envelope<Self>>,
 }
 
+impl GetNew for flume::Sender<Envelope<InMemory>> {
+  fn get_new(&self) -> Self { self.clone() }
+}
+
 impl Sender for flume::Sender<Envelope<InMemory>> {
   type Connection = InMemory;
 

arbiter-core/src/connection/mod.rs

@@ -1,4 +1,10 @@
-use std::{collections::HashMap, future::Future, hash::Hash, pin::Pin};
+use std::{
+  collections::HashMap,
+  future::Future,
+  hash::Hash,
+  pin::Pin,
+  sync::{Arc, Mutex},
+};
 
 use crate::handler::{Envelope, Message, Package};
 
@@ -26,8 +32,12 @@ impl Runtime for AsyncRuntime {
   fn wrap<T: 'static>(value: T) -> Self::Output<T> { Box::pin(async move { value }) }
 }
 
+pub trait GetNew {
+  fn get_new(&self) -> Self;
+}
+
 // TODO: Need to have results for send and receive.
-pub trait Sender: Send + Sync + 'static {
+pub trait Sender: GetNew + Send + Sync + 'static {
   type Connection: Connection;
   fn send(&self, envelope: Envelope<Self::Connection>);
 }
@@ -56,27 +66,35 @@ pub trait Connection: Send + Sync + 'static {
 
 pub struct Transport<C: Connection> {
   pub(crate) inbound_connection:   C,
-  pub(crate) outbound_connections: HashMap<C::Address, C::Sender>,
+  pub(crate) outbound_connections: Arc<Mutex<HashMap<C::Address, C::Sender>>>,
 }
 
 // TODO: These should return results
 impl<C: Connection> Transport<C> {
   pub fn new() -> Self {
-    Self { inbound_connection: C::new(), outbound_connections: HashMap::new() }
+    Self {
+      inbound_connection:   C::new(),
+      outbound_connections: Arc::new(Mutex::new(HashMap::new())),
+    }
   }
 
   pub fn send(&mut self, envelope: Envelope<C>, address: C::Address) {
-    self.outbound_connections.get_mut(&address).map(|sender| sender.send(envelope));
+    self.outbound_connections.lock().unwrap().get_mut(&address).map(|sender| sender.send(envelope));
   }
 
   pub fn broadcast(&mut self, envelope: Envelope<C>) {
-    self.outbound_connections.values_mut().for_each(|sender| sender.send(envelope.clone()));
+    self
+      .outbound_connections
+      .lock()
+      .unwrap()
+      .values_mut()
+      .for_each(|sender| sender.send(envelope.clone()));
   }
 
   pub fn receive(&mut self) -> Option<Envelope<C>> { self.inbound_connection.receiver().receive() }
 
   pub fn add_outbound_connection(&mut self, address: C::Address, sender: C::Sender) {
-    self.outbound_connections.insert(address, sender);
+    self.outbound_connections.lock().unwrap().insert(address, sender);
   }
 
   pub fn create_inbound_connection(&self) -> (C::Address, C::Sender) {

arbiter-core/src/connection/tcp.rs

@@ -6,14 +6,18 @@ use std::{
 };
 
 use crate::{
-  connection::{Connection, Receiver, Sender},
+  connection::{Connection, GetNew, Receiver, Sender},
   handler::Envelope,
 };
 
 pub struct Tcp {
   pub(crate) stream: TcpStream,
 }
 
+impl GetNew for TcpStream {
+  fn get_new(&self) -> Self { self.try_clone().unwrap() }
+}
+
 impl Sender for TcpStream {
   type Connection = Tcp;
 

arbiter-core/src/fabric.rs

@@ -1,8 +1,12 @@
-use std::collections::HashMap;
+use std::{
+  collections::HashMap,
+  sync::{Arc, Mutex},
+  thread::JoinHandle,
+};
 
 use crate::{
-  agent::{Agent, LifeCycle, RuntimeAgent, State},
-  connection::{memory::InMemory, Connection},
+  agent::{Agent, LifeCycle, RunningAgent, RuntimeAgent, State},
+  connection::{memory::InMemory, Connection, GetNew, Transport},
   handler::Package,
 };
 
@@ -11,7 +15,7 @@ use crate::{
 // and an atomic for their state)
 pub struct Fabric<C: Connection> {
   id:         FabricId,
-  agents:     HashMap<C::Address, Box<dyn RuntimeAgent<C>>>,
+  agents:     HashMap<C::Address, RunningAgent<C>>,
   name_to_id: HashMap<String, C::Address>,
 }
 
@@ -40,18 +44,20 @@ impl<C: Connection> Fabric<C> {
       }
       self.name_to_id.insert(name, new_agent.address());
     }
+
     let id = new_agent.address();
-    for agent in self.agents.values_mut() {
+    for (agent_address, agent) in self.agents.iter_mut() {
       // Get the new agent's inbound connection and give it to all the other agents as an outbound
       // connection.
-      let (address, sender) = new_agent.transport().create_inbound_connection();
-      agent.add_outbound_connection(address, sender);
+      let (address, sender) = new_agent.transport.create_inbound_connection();
+      agent.outbound_connections.lock().unwrap().insert(address, sender);
+
       // Get the other agent's inbound connection and give it to the new agent as an outbound
       // connection.
-      let (address, sender) = agent.transport().create_inbound_connection();
-      new_agent.add_outbound_connection(address, sender);
+      new_agent.add_outbound_connection(*agent_address, agent.sender.get_new());
     }
-    self.agents.insert(id, Box::new(new_agent));
+    let new_agent = new_agent.process();
+    self.agents.insert(id, new_agent);
     id
   }
 
@@ -65,7 +71,7 @@ impl<C: Connection> Fabric<C> {
     A: LifeCycle,
     C::Payload: Package<A::StartMessage> + Package<A::StopMessage>,
   {
-    agent.name = Some(name.into());
+    agent.set_name(name);
     Ok(self.register_agent(agent))
   }
 
@@ -104,7 +110,7 @@ impl<C: Connection> Fabric<C> {
     self.agents.get_mut(&agent_id).map_or_else(
       || Err(format!("Agent with ID {agent_id} not found")),
       |agent| {
-        agent.signal_start();
+        agent.controller.signal_start();
         Ok(())
       },
     )
@@ -124,7 +130,7 @@ found"
   // TODO: This is a bit clunky. Perhaps we can combine these states.
   /// Get the state of an agent by ID
   pub fn agent_state_by_id(&self, agent_id: C::Address) -> Option<State> {
-    self.agents.get(&agent_id).map(|agent| agent.current_loop_state())
+    self.agents.get(&agent_id).map(|agent| *agent.controller.state.lock().unwrap())
   }
 
   /// Get agent count
@@ -138,7 +144,7 @@ impl<C: Connection> Fabric<C> {
   /// Execute a single fabric step: poll transport and process messages
   pub fn start(&mut self) {
     for agent in self.agents.values_mut() {
-      agent.signal_start();
+      agent.controller.signal_start();
     }
   }
 }
@@ -230,7 +236,10 @@ mod tests {
       .unwrap();
 
     fabric.start();
-    let logger = fabric.agents.get(&logger_id).unwrap();
+    std::thread::sleep(std::time::Duration::from_millis(50));
+    let running_logger = fabric.agents.remove(&logger_id).unwrap();
+    running_logger.controller.signal_stop();
+    let logger = running_logger.task.join().unwrap();
     let logger = logger.inner_as_any().downcast_ref::<Logger>().unwrap();
     assert_eq!(logger.message_count, 1);
   }