Merge pull request #109 from 191220029/IP_management

Author: genedna

Cargo.toml

@@ -19,7 +19,7 @@ default-members = [".", "dagrs-derive"]
 
 [dependencies]
 dagrs-derive = { path = "dagrs-derive", optional = true, version = "0.4.3" }
-tokio = { version = "1.28", features = ["rt", "sync", "rt-multi-thread"] }
+tokio = { version = "1.28", features = ["rt", "sync", "rt-multi-thread", "time"] }
 log = "0.4"
 async-trait = "0.1.83"
 futures = "0.3.31"

examples/recv_any_example.rs

@@ -0,0 +1,187 @@
+//! # Example: recv_any_example
+//! Demonstrates how to use the `recv_any` method of `InChannels` to receive data from any available channel.
+//!
+//!
+//! # Output
+//! When running this example, you will see output similar to:
+//! ```
+//! Received message 'Hello from Sender' from node NodeId(1)
+//! Received message 'Hello from SlowSender' from node NodeId(2)
+//! ```
+//!
+//! The first message comes from the normal sender, and the second message comes from the slow sender
+//! after a 500ms delay.
+//!
+//! //! This example creates a graph with two senders and one receiver:
+//! - A normal sender that sends messages immediately
+//! - A slow sender that delays 500ms before sending messages
+//! - A receiver that uses `recv_any` to receive messages from either sender
+//!
+//! # Output
+//! When running this example, you will see output similar to:
+//! ```
+//! Received message 'Hello from Sender' from node NodeId(1)
+//! Received message 'Hello from SlowSender' from node NodeId(2)
+//! ```
+//!
+//! The first message comes from the normal sender, and the second message comes from the slow sender
+//! after a 500ms delay.
+
+use std::{sync::Arc, time::Duration};
+
+use async_trait::async_trait;
+use dagrs::{
+    connection::{in_channel::TypedInChannels, out_channel::TypedOutChannels},
+    node::typed_action::TypedAction,
+    DefaultNode, EnvVar, Graph, Node, NodeTable, Output,
+};
+use tokio::time::sleep;
+
+/// An action that sends a message to its output channel
+#[derive(Default)]
+pub struct SenderAction {
+    message: String,
+}
+
+impl SenderAction {
+    pub fn new(message: String) -> Self {
+        Self { message }
+    }
+}
+
+#[async_trait]
+impl TypedAction for SenderAction {
+    type I = ();
+    type O = String;
+
+    async fn run(
+        &self,
+        _: TypedInChannels<Self::I>,
+        out: TypedOutChannels<Self::O>,
+        _: Arc<EnvVar>,
+    ) -> Output {
+        // Send the message to all receivers
+        out.broadcast(self.message.clone()).await;
+        Output::Out(None)
+    }
+}
+
+/// An action that sends a message to its output channel after a delay
+#[derive(Default)]
+pub struct SlowSenderAction {
+    message: String,
+}
+
+impl SlowSenderAction {
+    pub fn new(message: String) -> Self {
+        Self { message }
+    }
+}
+
+#[async_trait]
+impl TypedAction for SlowSenderAction {
+    type I = ();
+    type O = String;
+
+    async fn run(
+        &self,
+        _: TypedInChannels<Self::I>,
+        out: TypedOutChannels<Self::O>,
+        _: Arc<EnvVar>,
+    ) -> Output {
+        // Wait for 500ms before sending
+        sleep(Duration::from_millis(500)).await;
+        // Send the message to all receivers
+        out.broadcast(self.message.clone()).await;
+        Output::Out(None)
+    }
+}
+
+/// An action that receives messages from any available channel
+#[derive(Default)]
+pub struct ReceiverAction;
+
+#[async_trait]
+impl TypedAction for ReceiverAction {
+    type I = String;
+    type O = ();
+
+    async fn run(
+        &self,
+        mut input: TypedInChannels<Self::I>,
+        _: TypedOutChannels<Self::O>,
+        _: Arc<EnvVar>,
+    ) -> Output {
+        // Receive from any available channel
+        match input.recv_any().await {
+            Ok((sender_id, content)) => {
+                let message = content.unwrap();
+                println!("Received message '{}' from node {:?}", message, sender_id);
+            }
+            Err(e) => {
+                eprintln!("Error receiving message: {:?}", e);
+            }
+        }
+
+        match input.recv_any().await {
+            Ok((sender_id, content)) => {
+                let message = content.unwrap();
+                println!("Received message '{}' from node {:?}", message, sender_id);
+            }
+            Err(e) => {
+                eprintln!("Error receiving message: {:?}", e);
+            }
+        }
+
+        Output::Out(None)
+    }
+}
+
+fn main() {
+    // Create a node table
+    let mut node_table = NodeTable::new();
+
+    // Create sender nodes
+    let sender1 = DefaultNode::with_action(
+        "Sender1".to_string(),
+        SenderAction::new("Hello from Sender".to_string()),
+        &mut node_table,
+    );
+    let sender2 = DefaultNode::with_action(
+        "Sender2".to_string(),
+        SlowSenderAction::new("Hello from SlowSender".to_string()),
+        &mut node_table,
+    );
+
+    // Create receiver node
+    let receiver = DefaultNode::with_action(
+        "Receiver".to_string(),
+        ReceiverAction::default(),
+        &mut node_table,
+    );
+
+    // Get node IDs before adding nodes to the graph
+    let sender1_id = sender1.id();
+    let sender2_id = sender2.id();
+    let receiver_id = receiver.id();
+
+    // Create a graph
+    let mut graph = Graph::new();
+
+    // Add nodes to the graph
+    graph.add_node(sender1);
+    graph.add_node(sender2);
+    graph.add_node(receiver);
+
+    // Add edges: both senders connect to the receiver
+    graph.add_edge(sender1_id, vec![receiver_id]);
+    graph.add_edge(sender2_id, vec![receiver_id]);
+
+    // Run the graph
+    match graph.start() {
+        Ok(_) => (),
+        Err(e) => {
+            eprintln!("Graph execution failed: {:?}", e);
+        }
+    }
+}

examples/typed_action.rs

@@ -116,7 +116,7 @@ fn main() {
             assert_eq!(*res, 272)
         }
         Err(e) => {
-            panic!("图执行失败: {:?}", e);
+            panic!("Graph execution failed: {:?}", e);
         }
     }
 }

src/connection/in_channel.rs

@@ -1,6 +1,7 @@
 use std::{collections::HashMap, marker::PhantomData, sync::Arc};
 
 use futures::future::join_all;
+use futures::future::select_ok;
 use tokio::sync::{broadcast, mpsc, Mutex};
 
 use crate::node::node::NodeId;
@@ -29,6 +30,31 @@ impl InChannels {
         }
     }
 
+    /// Receives data from any available channel and returns both the sender's ID and the content.
+    /// This method will wait until any channel has data available.
+    pub async fn recv_any(&mut self) -> Result<(NodeId, Content), RecvErr> {
+        let mut futures = Vec::new();
+        let ids: Vec<NodeId> = self.keys();
+
+        for id in ids {
+            let channel = self.get(&id).ok_or(RecvErr::NoSuchChannel)?;
+            let fut = Box::pin(async move {
+                let content = channel.lock().await.recv().await?;
+                Ok::<_, RecvErr>((id, content))
+            });
+            futures.push(fut);
+        }
+
+        if futures.is_empty() {
+            return Err(RecvErr::NoSuchChannel);
+        }
+
+        match select_ok(futures).await {
+            Ok((result, _)) => Ok(result),
+            Err(_) => Err(RecvErr::Closed),
+        }
+    }
+
     /// Calls `blocking_recv` for all the [`InChannel`]s, and applies transformation `f` to
     /// the return values of the call.
     pub fn blocking_map<F, T>(&mut self, mut f: F) -> Vec<T>
@@ -184,6 +210,31 @@ impl<T: Send + Sync + 'static> TypedInChannels<T> {
         }
     }
 
+    /// Receives typed data from any available channel and returns both the sender's ID and the typed content.
+    /// This method will wait until any channel has data available.
+    pub async fn recv_any(&mut self) -> Result<(NodeId, Option<Arc<T>>), RecvErr> {
+        let mut futures = Vec::new();
+        let ids: Vec<NodeId> = self.keys();
+
+        for id in ids {
+            let channel = self.get(&id).ok_or(RecvErr::NoSuchChannel)?;
+            let fut = Box::pin(async move {
+                let content: Content = channel.lock().await.recv().await?;
+                Ok::<_, RecvErr>((id, content.into_inner()))
+            });
+            futures.push(fut);
+        }
+
+        if futures.is_empty() {
+            return Err(RecvErr::NoSuchChannel);
+        }
+
+        match select_ok(futures).await {
+            Ok((result, _)) => Ok(result),
+            Err(_) => Err(RecvErr::Closed),
+        }
+    }
+
     /// Calls `blocking_recv` for all the [`InChannel`]s, and applies transformation `f` to
     /// the return values of the call.
     pub fn blocking_map<F, U>(&mut self, mut f: F) -> Vec<U>
@@ -225,14 +276,6 @@ impl<T: Send + Sync + 'static> TypedInChannels<T> {
         }
     }
 
-    pub(crate) fn insert(&mut self, node_id: NodeId, channel: Arc<Mutex<InChannel>>) {
-        self.0.insert(node_id, channel);
-    }
-
-    pub(crate) fn close_all(&mut self) {
-        self.0.values_mut().for_each(|c| c.blocking_lock().close());
-    }
-
     fn get(&self, id: &NodeId) -> Option<Arc<Mutex<InChannel>>> {
         match self.0.get(id) {
             Some(c) => Some(c.clone()),

src/connection/out_channel.rs

@@ -69,6 +69,11 @@ impl OutChannels {
     pub(crate) fn insert(&mut self, node_id: NodeId, channel: Arc<Mutex<OutChannel>>) {
         self.0.insert(node_id, channel);
     }
+
+    /// Returns a list of all available receiver node IDs.
+    pub fn get_receiver_ids(&self) -> Vec<NodeId> {
+        self.0.keys().copied().collect()
+    }
 }
 
 /// # Output Channel
@@ -179,18 +184,15 @@ impl<T: Send + Sync + 'static> TypedOutChannels<T> {
         }
     }
 
-    pub(crate) fn close_all(&mut self) {
-        self.0.clear();
-    }
-
     fn get(&self, id: &NodeId) -> Option<Arc<Mutex<OutChannel>>> {
         match self.0.get(id) {
             Some(c) => Some(c.clone()),
             None => None,
         }
     }
 
-    pub(crate) fn insert(&mut self, node_id: NodeId, channel: Arc<Mutex<OutChannel>>) {
-        self.0.insert(node_id, channel);
+    /// Returns a list of all available receiver node IDs.
+    pub fn get_receiver_ids(&self) -> Vec<NodeId> {
+        self.0.keys().copied().collect()
     }
 }