Add timeouts for disconnects. (#17)

Author: doru1004

examples/http/http_source.py

@@ -16,7 +16,6 @@
 
 import ray
 import rayvens
-import time
 import sys
 
 # Receive message from stock price source and print it to console using the
@@ -27,7 +26,7 @@
     print(f'usage: {sys.argv[0]} <run_mode>')
     sys.exit(1)
 run_mode = sys.argv[1]
-if run_mode not in ['operator']:
+if run_mode not in ['local', 'operator']:
     raise RuntimeError(f'Invalid run mode provided: {run_mode}')
 
 # Initialize ray either on the cluster or locally otherwise.
@@ -67,7 +66,5 @@
 # Log all events from stream-attached sources.
 stream >> (lambda event: print('LOG:', event))
 
-# Wait before ending program.
-time.sleep(20)
-
-stream.disconnect_all()
+# Disconnect all sources after 20 seconds.
+stream.disconnect_all(after=20)

examples/slack/slack_operator_mode.py

@@ -17,7 +17,6 @@
 import ray
 import rayvens
 import sys
-import time
 
 # Send message to Slack sink using the kamel anywhere operator implementation.
 
@@ -50,14 +49,11 @@
                    webhookUrl=slack_webhook)
 
 # Add sink to stream.
-sink = stream.actor.add_sink.remote(stream, sink_config)
-
-# Wait for sink to reach running state.
-ray.get(sink)
+sink = stream.add_sink(sink_config)
 
 # Sends message to all sinks attached to this stream.
 stream << f'Sending message to Slack sink in run mode {run_mode}.'
 
-time.sleep(5)
-
-ray.get(stream.actor.disconnect_all.remote())
+# Disconnect any sources or sinks attached to the stream 2 seconds after
+# the stream is idle (i.e. no events were propagated by the stream).
+stream.disconnect_all(after_idle_for=2)

rayvens/api.py

@@ -16,6 +16,7 @@
 
 import os
 import ray
+import time
 
 from rayvens.core.local import start as start_http
 from rayvens.core.kafka import start as start_kafka
@@ -63,18 +64,38 @@ def add_sink(self, sink_config):
     def unsubscribe(self, subscriber_name):
         return ray.get(self.actor.unsubscribe.remote(subscriber_name))
 
-    def disconnect_source(self, source_name):
+    def disconnect_source(self, source_name, after_idle_for=None, after=None):
+        self._wait_for_timeout(after_idle_for)
         return ray.get(self.actor.disconnect_source.remote(source_name))
 
-    def disconnect_sink(self, sink_name):
+    def disconnect_sink(self, sink_name, after_idle_for=None, after=None):
+        self._wait_for_timeout(after_idle_for)
         return ray.get(self.actor.disconnect_sink.remote(sink_name))
 
-    def disconnect_all(self):
+    def disconnect_all(self, after_idle_for=None, after=None):
+        self._wait_for_timeout(after_idle_for, after)
         return ray.get(self.actor.disconnect_all.remote())
 
     def _meta(self, action, *args, **kwargs):
         return ray.get(self.actor._meta.remote(action, *args, **kwargs))
 
+    def _wait_for_timeout(self, after_idle_for, after):
+        if after_idle_for is not None and after_idle_for > 0:
+            while True:
+                time_elapsed_since_last_event = self._idle_time()
+
+                # Idle timeout exceeds the user-specified time limit:
+                if time_elapsed_since_last_event > after_idle_for:
+                    break
+
+                # Check again after waiting for the rest of the timeout time:
+                time.sleep(after_idle_for - time_elapsed_since_last_event + 1)
+        if after is not None and after > 0:
+            time.sleep(after)
+
+    def _idle_time(self):
+        return time.time() - ray.get(self.actor._get_latest_timestamp.remote())
+
 
 @ray.remote(num_cpus=0)
 class StreamActor:
@@ -84,6 +105,7 @@ def __init__(self, name, operator=None):
         self._operator = operator
         self._sources = {}
         self._sinks = {}
+        self._latest_sent_event_timestamp = None
 
     def send_to(self, subscriber, name=None):
         if name in self._subscribers:
@@ -104,6 +126,7 @@ def append(self, data):
                 if not integration.accepts_data_type(data):
                     continue
             _eval(subscriber, data)
+        self._latest_sent_event_timestamp = time.time()
 
     def add_operator(self, operator):
         self._operator = operator
@@ -159,6 +182,9 @@ def disconnect_all(self):
     def _meta(self, action, *args, **kwargs):
         return verify_do(self, _global_camel, action, *args, **kwargs)
 
+    def _get_latest_timestamp(self):
+        return self._latest_sent_event_timestamp
+
 
 def _eval(f, data):
     if isinstance(f, Stream):