1 files changed, 17 insertions, 29 deletions

diff --git a/crates/core/tedge_api/examples/heartbeat.rs b/crates/core/tedge_api/examples/heartbeat.rs
index 95e6ca15..37e9c8ec 100644
--- a/crates/core/tedge_api/examples/heartbeat.rs
+++ b/crates/core/tedge_api/examples/heartbeat.rs
@@ -1,15 +1,11 @@
-use std::{
-    any::TypeId,
-    collections::{HashMap, HashSet},
-    time::Duration,
-};
+use std::{collections::HashMap, time::Duration};
 
 use async_trait::async_trait;
 use futures::FutureExt;
 use tedge_api::{
-    address::ReplySender,
-    message::NoReply,
-    plugin::{BuiltPlugin, Handle, Message, PluginDeclaration, PluginExt},
+    address::ReplySenderFor,
+    message::MessageType,
+    plugin::{AcceptsReplies, BuiltPlugin, Handle, Message, PluginDeclaration, PluginExt},
     Address, CancellationToken, Plugin, PluginBuilder, PluginConfiguration, PluginDirectory,
     PluginError,
 };
@@ -17,7 +13,8 @@ use tedge_api::{
 /// A message that represents a heartbeat that gets sent to plugins
 #[derive(Debug)]
 struct Heartbeat;
-impl Message for Heartbeat {
+impl Message for Heartbeat {}
+impl AcceptsReplies for Heartbeat {
     type Reply = HeartbeatStatus;
 }
 
@@ -27,9 +24,7 @@ enum HeartbeatStatus {
     Alive,
     Degraded,
 }
-impl Message for HeartbeatStatus {
-    type Reply = NoReply;
-}
+impl Message for HeartbeatStatus {}
 
 /// A PluginBuilder that gets used to build a HeartbeatService plugin instance
 #[derive(Debug)]
@@ -253,7 +248,7 @@ impl Handle<Heartbeat> for CriticalService {
     async fn handle_message(
         &self,
         _message: Heartbeat,
-        sender: ReplySender<HeartbeatStatus>,
+        sender: ReplySenderFor<Heartbeat>,
     ) -> Result<(), PluginError> {
         println!("CriticalService: Received Heartbeat!");
         let mut status = self.status.lock().await;
@@ -301,23 +296,13 @@ impl Plugin for CriticalService {
 /// Helper type for keeping information about plugins during runtime
 #[derive(Debug)]
 struct PluginInfo {
-    types: HashSet<(&'static str, TypeId)>,
+    types: Vec<MessageType>,
     receiver: Option<tedge_api::address::MessageReceiver>,
     sender: tedge_api::address::MessageSender,
 }
 
-impl Clone for PluginInfo {
-    fn clone(&self) -> Self {
-        PluginInfo {
-            types: self.types.clone(),
-            receiver: None,
-            sender: self.sender.clone(),
-        }
-    }
-}
-
 /// The type that provides the communication infrastructure to the plugins.
-#[derive(Clone, Debug)]
+#[derive(Debug)]
 struct Communication {
     plugins: HashMap<String, PluginInfo>,
 }
@@ -328,7 +313,7 @@ impl Communication {
         self.plugins.insert(
             name.to_owned(),
             PluginInfo {
-                types: PB::kind_message_types().into(),
+                types: PB::kind_message_types().into_types(),
                 sender,
                 receiver: Some(receiver),
             },
@@ -341,7 +326,7 @@ impl PluginDirectory for Communication {
         &self,
         name: &str,
     ) -> Result<Address<MB>, tedge_api::error::DirectoryError> {
-        let types = MB::get_ids().into_iter().collect();
+        let asked_types: Vec<_> = MB::get_ids().into_iter().collect();
 
         let plug = self.plugins.get(name).unwrap_or_else(|| {
             // This is an example, so we panic!() here.
@@ -353,12 +338,15 @@ impl PluginDirectory for Communication {
             )
         });
 
-        if !plug.types.is_superset(&types) {
+        if !asked_types
+            .iter()
+            .all(|req_type| plug.types.iter().any(|ty| ty.satisfy(req_type)))
+        {
             // This is an example, so we panic!() here
             // In real-world, we would do some reporting and return an error
             panic!(
                 "Asked for {:#?} but plugin {} only has types {:#?}",
-                types, name, plug.types,
+                asked_types, name, plug.types,
             );
         } else {
             Ok(Address::new(plug.sender.clone()))