Cleaning plugin management code

2 files changed, 11 insertions, 14 deletions

diff --git a/esds/node.py b/esds/node.py
index f4e1d16..238c08d 100644
--- a/esds/node.py
+++ b/esds/node.py
@@ -22,7 +22,7 @@ class Node:
     def plugin_register(self,plugin):
         self.plugins.append(plugin)
 
-    def plugin_notify(self,reason,args):
+    def plugin_notify(self,reason,args,time=None):
         """
         This function strives to avoid using Python specific features
         """
@@ -34,7 +34,7 @@ class Node:
             if reason == "send_return":
                 p.on_send_return(args[0],args[1],args[2],args[3],args[4])
             if reason == "on_communication_end":
-                p.on_communication_end(args[0],args[1])
+                p.on_communication_end(args[0],time,args[1])
             if reason == "terminated":
                 p.on_terminated()
 
@@ -202,7 +202,7 @@ class Node:
         while True:
             ack=self.rqueue.get() # Wait for simulator acknowledgments
             if ack[0] == "plugin_notify":
-                self.plugin_notify(ack[1],ack[2])
+                self.plugin_notify(ack[1],ack[3],time=ack[2])
                 self["pending_plugin_notify"]-=1
             elif ack[0] not in ack_types:
                 ack_buffer.append(ack)
@@ -217,7 +217,7 @@ class Node:
         """
         Wait until node stop running
         """
-        while self["state"] == "running":
+        while self["state"] == "running" or self["pending_plugin_notify"] > 0:
             pass
         
     def run(self,args):
diff --git a/esds/simulator.py b/esds/simulator.py
index aa6573b..4f2eb71 100644
--- a/esds/simulator.py
+++ b/esds/simulator.py
@@ -202,7 +202,7 @@ class Simulator:
                             else:
                                 selector_wireless.append(True)
                                 selector_wired.append(False)
-                            self.notify_node_plugins(self.nodes[int(event[2][1])], "on_communication_end", (self.time,event))
+                            self.notify_node_plugins(self.nodes[int(event[2][1])], "on_communication_end", event)
                         else:
                             selector_wireless.append(False)
                             selector_wired.append(False)
@@ -235,7 +235,7 @@ class Simulator:
                             selector.append(True)
                             if self.netmat[event[2][2]]["is_wired"]:
                                 sharing_to_update.append((int(event[2][1]),event[2][2]))
-                            self.notify_node_plugins(node, "on_communication_end", (self.time,event))
+                            self.notify_node_plugins(node, "on_communication_end", event)
                         else:
                             selector.append(False)
                     self.events=self.events[~np.array(selector)]
@@ -359,10 +359,7 @@ class Simulator:
 
     def notify_node_plugins(self,node,callback,args):
         node["pending_plugin_notify"]+=1
-        node.rqueue.put(("plugin_notify",callback,args))
-        # Now ensure that all callbacks are called before continuing
-        while node["pending_plugin_notify"] > 0:
-            pass
+        node.rqueue.put(("plugin_notify",callback,self.time,args))
 
     def add_event(self,event_type,event_ts,event,priority=2):
         """
@@ -438,13 +435,13 @@ class Simulator:
                                 dst["state"]="running"
                                 dst.rqueue.put(("receive",0))
                                 self.sync_node_non_blocking(dst,timeout_remove_only=True)
-                            self.notify_node_plugins(dst, "on_communication_end", (self.time,event))
+                            self.notify_node_plugins(dst, "on_communication_end", event)
                         self.update_sharing(dst.node_id,-1,interface)
                         src["state"]="running"
                         code=0 if perform_delivery else 1
                         src.rqueue.put(("send",code))
                         self.sync_node_non_blocking(src,timeout_remove_only=True)
-                        self.notify_node_plugins(src, "on_communication_end", (self.time,event))
+                        self.notify_node_plugins(src, "on_communication_end", event)
                     else:
                         if src.node_id != dst.node_id:
                             if perform_delivery:
@@ -457,12 +454,12 @@ class Simulator:
                                     dst["state"]="running"
                                     dst.rqueue.put(("receive",0))
                                     self.sync_node_non_blocking(dst,timeout_remove_only=True)
-                                self.notify_node_plugins(dst, "on_communication_end", (self.time,event))
+                                self.notify_node_plugins(dst, "on_communication_end", event)
                         else:
                             src["state"]="running"
                             src.rqueue.put(("send",0))
                             self.sync_node_non_blocking(src,timeout_remove_only=True)
-                            self.notify_node_plugins(src, "on_communication_end", (self.time,event))
+                            self.notify_node_plugins(src, "on_communication_end", event)
                 elif event_type == 1: # Timeout
                     node=self.nodes[int(event)]
                     node["state"]="running"