Gracefully end source-related threads (#53)

doru1004 · web-flow · commit 00ac46e68198 · 2021-09-03T13:39:11.000-04:00
* Gracefully end source-related threads.

* Add timer for waiting for requests before joining thread.
diff --git a/rayvens/core/common.py b/rayvens/core/common.py
@@ -125,22 +125,37 @@ def append(self, data):
             pass
 
 
-# Method for sending events to the queue.
-def send_to(handle, server_address, route):
-    def append():
-        while True:
+class QueueReadThread(threading.Thread):
+    def __init__(self, handle, server_address, route):
+        threading.Thread.__init__(self)
+        self.stop_flag = threading.Event()
+        self.run_flag = threading.Event()
+        self.handle = handle
+        self.server_address = server_address
+        self.route = route
+
+    def run(self):
+        while not self.stop_flag.is_set():
             try:
-                response = requests.get(f'{server_address}{route}',
+                response = requests.get(f'{self.server_address}{self.route}',
                                         timeout=(5, None))
                 if response.status_code != 200:
                     time.sleep(1)
                     continue
                 response.encoding = 'latin-1'
-                handle.append.remote(response.text)
+                self.handle.append.remote(response.text)
             except requests.exceptions.ConnectionError:
                 time.sleep(1)
 
-    threading.Thread(target=append).start()
+        self.run_flag.set()
+
+
+# Method for sending events to the queue.
+def send_to(handle, server_address, route):
+    queue_reading_thread = QueueReadThread(handle, server_address, route)
+    queue_reading_thread.daemon = True
+    queue_reading_thread.start()
+    return queue_reading_thread
 
 
 # Method for sending events to the sink.
@@ -222,30 +237,63 @@ def close_consumer(self):
         self.kafka_consumer.close()
 
 
+class ConsumerEnablingThread(threading.Thread):
+    def __init__(self, handle, consumers):
+        threading.Thread.__init__(self)
+        self.stop_flag = threading.Event()
+        self.run_flag = threading.Event()
+        self.handle = handle
+        self.consumers = consumers
+
+    def run(self):
+        subscribers, operator = ray.get(self.handle._fetch_processors.remote())
+        try:
+            consumer_references = [
+                consumer.enable_consumer.remote(subscribers, operator)
+                for consumer in self.consumers
+            ]
+            ray.get(consumer_references)
+        except KeyboardInterrupt:
+            for consumer in self.consumers:
+                consumer.disable_consumer.remote()
+        finally:
+            for consumer in self.consumers:
+                consumer.close_consumer.remote()
+
+        self.run_flag.set()
+
+
+class ConsumerThread(threading.Thread):
+    def __init__(self, handle, consumer):
+        threading.Thread.__init__(self)
+        self.stop_flag = threading.Event()
+        self.run_flag = threading.Event()
+        self.handle = handle
+        self.consumer = consumer
+
+    def run(self):
+        while not self.stop_flag.is_set():
+            msg = self.consumer.poll()
+            if msg.error():
+                print(f'consumer error: {msg.error()}')
+            else:
+                self.handle.append.remote(msg.value().decode('utf-8'))
+
+        self.run_flag.set()
+
+
 def kafka_send_to(kafka_transport_topic, kafka_transport_partitions, handle):
+    # If source uses a scaling transport start multiple Kafka Consumers.
     if kafka_transport_partitions > 1:
         consumers = [
             KafkaConsumer.remote(kafka_transport_topic, handle)
             for _ in range(kafka_transport_partitions)
         ]
 
-        def append():
-            subscribers, operator = ray.get(handle._fetch_processors.remote())
-            try:
-                consumer_references = [
-                    consumer.enable_consumer.remote(subscribers, operator)
-                    for consumer in consumers
-                ]
-                ray.get(consumer_references)
-            except KeyboardInterrupt:
-                for consumer in consumers:
-                    consumer.disable_consumer.remote()
-            finally:
-                for consumer in consumers:
-                    consumer.close_consumer.remote()
-
-        threading.Thread(target=append).start()
-        return
+        consumer_enabling_thread = ConsumerEnablingThread(handle, consumers)
+        consumer_enabling_thread.daemon = True
+        consumer_enabling_thread.start()
+        return consumer_enabling_thread
 
     # Use kafka consumer thread to push from camel source to rayvens stream
     consumer = Consumer({
@@ -256,15 +304,10 @@ def append():
 
     consumer.subscribe([kafka_transport_topic])
 
-    def append():
-        while True:
-            msg = consumer.poll()
-            if msg.error():
-                print(f'consumer error: {msg.error()}')
-            else:
-                handle.append.remote(msg.value().decode('utf-8'))
-
-    threading.Thread(target=append).start()
+    consumer_thread = ConsumerThread(handle, consumer)
+    consumer_thread.daemon = True
+    consumer_thread.start()
+    return consumer_thread
 
 
 def kafka_recv_from(integration_name, kafka_transport_topic, handle):
diff --git a/rayvens/core/integration.py b/rayvens/core/integration.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 #
 
+import time
 from rayvens.core.common import brokers
 from rayvens.core.utils import random_port, create_partitioned_topic
 from rayvens.core.name import name_integration
@@ -58,6 +59,7 @@ def __init__(self, stream_name, source_sink_name, config):
         self.service_name = None
         self.server_address = None
         self.environment_preparators = []
+        self.thread = None
 
     def invoke_local_run(self, mode, integration_content):
         self.port = random_port(mode.check_port)
@@ -87,6 +89,27 @@ def invoke_run(self, mode, integration_content):
     # created using the kamel operator and `kamel run` or an integration
     # created using `kamel local run`.
     def disconnect(self, mode):
+        # Set stop_flag to true.
+        if self.thread is not None:
+            # Signal the end of ongoing loops.
+            self.thread.stop_flag.set()
+
+            # Wait for threads to perform perform/finish the current requests.
+            # Small 10 s timer to wait for request completion.
+            countdown = 10
+            while not self.thread.run_flag.is_set():
+                time.sleep(1)
+                if countdown == 0:
+                    break
+                countdown -= 1
+
+            # Join back successful threads otherwise threads will be killed
+            # when the application ends.
+            if self.thread.run_flag.is_set():
+                self.thread.join()
+                print("Successfully terminated additional"
+                      f" {self.integration_name} integration thread.")
+
         if self.invocation.uses_operator():
             # If kamel is running the cluster then use kamel delete to
             # terminate the integration. First we terminate any services
diff --git a/rayvens/core/kafka.py b/rayvens/core/kafka.py
@@ -36,8 +36,9 @@ def add_source(self, stream, config, source_name):
             f'kafka:{integration.kafka_transport_topic}?brokers={brokers()}')
         integration.prepare_environment(self.mode)
         integration.invoke_local_run(self.mode, spec)
-        kafka_send_to(integration.kafka_transport_topic,
-                      integration.kafka_transport_partitions, stream.actor)
+        integration.thread = kafka_send_to(
+            integration.kafka_transport_topic,
+            integration.kafka_transport_partitions, stream.actor)
         return integration
 
     def add_sink(self, stream, config, sink_name):
diff --git a/rayvens/core/local.py b/rayvens/core/local.py
@@ -36,7 +36,9 @@ def add_source(self, stream, config, source_name):
                                         inverted=True)
         integration.prepare_environment(self.mode)
         integration.invoke_local_run(self.mode, spec)
-        send_to(stream.actor, self.mode.server_address(integration), route)
+        integration.thread = send_to(stream.actor,
+                                     self.mode.server_address(integration),
+                                     route)
         if not await_start(self.mode, integration):
             raise RuntimeError('Could not start source')
         return integration
diff --git a/rayvens/core/operator.py b/rayvens/core/operator.py
@@ -50,7 +50,9 @@ def add_source(self, stream, source, source_name):
         integration.invoke_run(self.mode, integration_content)
 
         # Set up source for the HTTP connector case.
-        send_to(stream.actor, self.mode.server_address(integration), route)
+        integration.thread = send_to(stream.actor,
+                                     self.mode.server_address(integration),
+                                     route)
 
         if not await_start(self.mode, integration):
             raise RuntimeError('Could not start source')
diff --git a/rayvens/core/operator_kafka.py b/rayvens/core/operator_kafka.py
@@ -47,8 +47,9 @@ def add_source(self, stream, source, source_name):
         integration.invoke_run(self.mode, integration_content)
 
         # Set up source for the HTTP connector case.
-        kafka_send_to(integration.kafka_transport_topic,
-                      integration.kafka_transport_partitions, stream.actor)
+        integration.thread = kafka_send_to(
+            integration.kafka_transport_topic,
+            integration.kafka_transport_partitions, stream.actor)
 
         if not await_start(self.mode, integration):
             raise RuntimeError('Could not start source')
