Browse Source

net/message_subscriber: new fn receive_with_timeout added on message subscription

skoupidi 2 years ago
parent
commit
b468b9c399
1 changed files with 16 additions and 2 deletions
  1. 16 2
      src/net/message_subscriber.rs

+ 16 - 2
src/net/message_subscriber.rs

@@ -16,7 +16,7 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
  */
 
 
-use std::{any::Any, collections::HashMap, io::Cursor, sync::Arc};
+use std::{any::Any, collections::HashMap, io::Cursor, sync::Arc, time::Duration};
 
 
 use async_trait::async_trait;
 use async_trait::async_trait;
 use futures::stream::{FuturesUnordered, StreamExt};
 use futures::stream::{FuturesUnordered, StreamExt};
@@ -25,7 +25,7 @@ use rand::{rngs::OsRng, Rng};
 use smol::lock::Mutex;
 use smol::lock::Mutex;
 
 
 use super::message::Message;
 use super::message::Message;
-use crate::{Error, Result};
+use crate::{system::timeout::timeout, Error, Result};
 
 
 /// 64-bit identifier for message subscription.
 /// 64-bit identifier for message subscription.
 pub type MessageSubscriptionId = u64;
 pub type MessageSubscriptionId = u64;
@@ -144,6 +144,20 @@ impl<M: Message> MessageSubscription<M> {
         }
         }
     }
     }
 
 
+    /// Start receiving messages with timeout.
+    pub async fn receive_with_timeout(&self, seconds: u64) -> MessageResult<M> {
+        let dur = Duration::from_secs(seconds);
+        let Ok(res) = timeout(dur, self.recv_queue.recv()).await else {
+            return Err(Error::ConnectTimeout)
+        };
+        match res {
+            Ok(message) => message,
+            Err(e) => {
+                panic!("MessageSubscription::receive_with_timeout(): recv_queue failed! {}", e)
+            }
+        }
+    }
+
     /// Unsubscribe from a message subscription. Must be called manually.
     /// Unsubscribe from a message subscription. Must be called manually.
     pub async fn unsubscribe(&self) {
     pub async fn unsubscribe(&self) {
         self.parent.unsubscribe(self.id).await
         self.parent.unsubscribe(self.id).await