feat(flow): block() now waits for specific threading event (#3790)

* feat(flow): block() now waits for specific threading event

* style: fix overload and cli autocomplete

Co-authored-by: Jina Dev Bot <dev-bot@jina.ai>

Author: hanxiao

docs/fundamentals/flow/flow-api.md

@@ -39,6 +39,28 @@ prevent that, use `.block()` to suspend the current process.
 with f:
     f.block()  # block the current process
 ```
+
+To terminate a blocked Flow, one can send use `threading.Event` or `multiprocessing.Event` to send the terminate signal.
+The following example terminates a Flow from another process. 
+
+```python
+import multiprocessing
+import time
+
+from jina import Flow
+
+ev = multiprocessing.Event()
+
+def close_after_5s():
+    time.sleep(5)
+    ev.set()
+
+f = Flow().add()
+with f:
+    t = multiprocessing.Process(target=close_after_5s)
+    t.start()
+    f.block(stop_event=ev)
+```
 ````
 ## Visualize a Flow
 

jina/flow/base.py

@@ -3,6 +3,7 @@
 import copy
 import itertools
 import json
+import multiprocessing
 import os
 import re
 import sys
@@ -1117,6 +1118,9 @@ def __enter__(self):
             return self.start()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        if hasattr(self, '_stop_event'):
+            self._stop_event.set()
+
         super().__exit__(exc_type, exc_val, exc_tb)
 
         # unset all envs to avoid any side-effect
@@ -1166,7 +1170,7 @@ def start(self):
 
         return self
 
-    def _wait_until_all_ready(self) -> bool:
+    def _wait_until_all_ready(self):
         results = {}
         threads = []
 
@@ -1578,10 +1582,22 @@ def _get_address_table(self, address_table):
             )
         return address_table
 
-    def block(self):
-        """Block the process until user hits KeyboardInterrupt"""
+    def block(
+        self, stop_event: Optional[Union[threading.Event, multiprocessing.Event]] = None
+    ):
+        """Block the Flow until `stop_event` is set or user hits KeyboardInterrupt
+
+        :param stop_event: a threading event or a multiprocessing event that onces set will resume the control Flow
+            to main thread.
+        """
         try:
-            threading.Event().wait()
+            if stop_event is None:
+                self._stop_event = (
+                    threading.Event()
+                )  #: this allows `.close` to close the Flow from another thread/proc
+                self._stop_event.wait()
+            else:
+                stop_event.wait()
         except KeyboardInterrupt:
             pass
 

jina/types/document/__init__.py

@@ -161,7 +161,7 @@ def __init__(
     ):
         """
         :param adjacency: the adjacency of this Document
-        :param blob: the blob content of this Document
+        :param blob: the blob content of thi Document
         :param buffer: the buffer bytes from this document
         :param chunks: the array of chunks of this document
         :param content: the value of the content depending on `:meth:`content_type`

tests/unit/flow-construct/test_flow_block.py

@@ -0,0 +1,45 @@
+import multiprocessing
+import threading
+import time
+
+from jina import Flow
+
+
+def test_closing_blocked_flow_from_another_thread_via_flow():
+    def close_blocked_f(flow):
+        time.sleep(3)
+        flow.close()
+
+    f = Flow().add()
+    with f:
+        t = threading.Thread(target=close_blocked_f, args=(f,))
+        t.start()
+        f.block()
+
+
+def test_closing_blocked_flow_from_another_thread_via_event():
+    ev = threading.Event()
+
+    def close_blocked_f():
+        time.sleep(3)
+        ev.set()
+
+    f = Flow().add()
+    with f:
+        t = threading.Thread(target=close_blocked_f)
+        t.start()
+        f.block(stop_event=ev)
+
+
+def test_closing_blocked_flow_from_another_process_via_event():
+    ev = multiprocessing.Event()
+
+    def close_blocked_f():
+        time.sleep(3)
+        ev.set()
+
+    f = Flow().add()
+    with f:
+        t = multiprocessing.Process(target=close_blocked_f)
+        t.start()
+        f.block(stop_event=ev)