Improve UDF cancellation

Author: kesmit13

singlestoredb/__init__.py

@@ -25,7 +25,7 @@
     DataError, ManagementError,
 )
 from .management import (
-    manage_cluster, manage_workspaces, manage_files,
+    manage_cluster, manage_workspaces, manage_files, manage_regions,
 )
 from .types import (
     Date, Time, Timestamp, DateFromTicks, TimeFromTicks, TimestampFromTicks,

singlestoredb/functions/ext/asgi.py

@@ -285,7 +285,12 @@ def cancel_on_event(
 
     """
     if cancel_event.is_set():
-        raise asyncio.CancelledError('Function call was cancelled')
+        task = asyncio.current_task()
+        if task is not None:
+            task.cancel()
+        raise asyncio.CancelledError(
+            'Function call was cancelled by client',
+        )
 
 
 def build_udf_endpoint(
@@ -314,19 +319,21 @@ def build_udf_endpoint(
 
         async def do_func(
             cancel_event: threading.Event,
+            finished_event: threading.Event,
             timer: Timer,
             row_ids: Sequence[int],
             rows: Sequence[Sequence[Any]],
         ) -> Tuple[Sequence[int], List[Tuple[Any, ...]]]:
             '''Call function on given rows of data.'''
             out = []
-            with timer('call_function'):
+            async with timer('call_function'):
                 for row in rows:
                     cancel_on_event(cancel_event)
                     if is_async:
                         out.append(await func(*row))
                     else:
                         out.append(func(*row))
+            finished_event.set()
             return row_ids, list(zip(out))
 
         return do_func
@@ -360,6 +367,7 @@ def build_vector_udf_endpoint(
 
     async def do_func(
         cancel_event: threading.Event,
+        finished_event: threading.Event,
         timer: Timer,
         row_ids: Sequence[int],
         cols: Sequence[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
@@ -371,7 +379,7 @@ async def do_func(
         row_ids = array_cls(row_ids)
 
         # Call the function with `cols` as the function parameters
-        with timer('call_function'):
+        async with timer('call_function'):
             if cols and cols[0]:
                 if is_async:
                     out = await func(*[x if m else x[0] for x, m in zip(cols, masks)])
@@ -383,6 +391,7 @@ async def do_func(
                 else:
                     out = func()
 
+        finished_event.set()
         cancel_on_event(cancel_event)
 
         # Single masked value
@@ -425,6 +434,7 @@ def build_tvf_endpoint(
 
         async def do_func(
             cancel_event: threading.Event,
+            finished_event: threading.Event,
             timer: Timer,
             row_ids: Sequence[int],
             rows: Sequence[Sequence[Any]],
@@ -433,7 +443,7 @@ async def do_func(
             out_ids: List[int] = []
             out = []
             # Call function on each row of data
-            with timer('call_function'):
+            async with timer('call_function'):
                 for i, row in zip(row_ids, rows):
                     cancel_on_event(cancel_event)
                     if is_async:
@@ -442,6 +452,7 @@ async def do_func(
                         res = func(*row)
                     out.extend(as_list_of_tuples(res))
                     out_ids.extend([row_ids[i]] * (len(out)-len(out_ids)))
+            finished_event.set()
             return out_ids, out
 
         return do_func
@@ -474,6 +485,7 @@ def build_vector_tvf_endpoint(
 
     async def do_func(
         cancel_event: threading.Event,
+        finished_event: threading.Event,
         timer: Timer,
         row_ids: Sequence[int],
         cols: Sequence[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
@@ -489,7 +501,7 @@ async def do_func(
         is_async = asyncio.iscoroutinefunction(func)
 
         # Call function on each column of data
-        with timer('call_function'):
+        async with timer('call_function'):
             if cols and cols[0]:
                 if is_async:
                     func_res = await func(
@@ -505,6 +517,8 @@ async def do_func(
                 else:
                     func_res = func()
 
+        finished_event.set()
+
         res = get_dataframe_columns(func_res)
 
         cancel_on_event(cancel_event)
@@ -616,15 +630,11 @@ async def cancel_on_disconnect(
             )
 
 
-def cancel_all_tasks(tasks: Iterable[asyncio.Task[Any]]) -> None:
+async def cancel_all_tasks(tasks: Iterable[asyncio.Task[Any]]) -> None:
     """Cancel all tasks."""
     for task in tasks:
-        if task.done():
-            continue
-        try:
-            task.cancel()
-        except Exception:
-            pass
+        task.cancel()
+    await asyncio.gather(*tasks, return_exceptions=True)
 
 
 def start_counter() -> float:
@@ -1027,17 +1037,24 @@ async def __call__(
                 result = []
 
                 cancel_event = threading.Event()
+                finished_event = threading.Event()
+
+                # Async functions don't need to set the finished event
+                if func_info['is_async']:
+                    finished_event.set()
 
                 with timer('parse_input'):
                     inputs = input_handler['load'](  # type: ignore
                         func_info['colspec'], b''.join(data),
                     )
 
                 func_task = asyncio.create_task(
-                    func(cancel_event, timer, *inputs)
+                    func(cancel_event, finished_event, timer, *inputs)
                     if func_info['is_async']
                     else to_thread(
-                        lambda: asyncio.run(func(cancel_event, timer, *inputs)),
+                        lambda: asyncio.run(
+                            func(cancel_event, finished_event, timer, *inputs),
+                        ),
                     ),
                 )
                 disconnect_task = asyncio.create_task(
@@ -1049,12 +1066,15 @@ async def __call__(
 
                 all_tasks += [func_task, disconnect_task, timeout_task]
 
-                with timer('function_wrapper'):
+                async with timer('function_wrapper'):
                     done, pending = await asyncio.wait(
                         all_tasks, return_when=asyncio.FIRST_COMPLETED,
                     )
 
-                cancel_all_tasks(pending)
+                await cancel_all_tasks(pending)
+
+                # Make sure threads finish before we proceed
+                finished_event.wait()
 
                 for task in done:
                     if task is disconnect_task:
@@ -1105,7 +1125,7 @@ async def __call__(
                 await send(self.error_response_dict)
 
             finally:
-                cancel_all_tasks(all_tasks)
+                await cancel_all_tasks(all_tasks)
 
         # Handle api reflection
         elif method == 'GET' and path == self.show_create_function_path:

singlestoredb/functions/ext/timer.py

@@ -111,6 +111,14 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         self.metrics.setdefault(timing_info['key'], 0)
         self.metrics[timing_info['key']] += elapsed
 
+    async def __aenter__(self) -> 'Timer':
+        """Async enter for async context manager support."""
+        return self.__enter__()
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Async exit for async context manager support."""
+        self.__exit__(exc_type, exc_val, exc_tb)
+
     def finish(self) -> None:
         """Finish the current timing context and store the elapsed time."""
         if self._stack:

singlestoredb/management/__init__.py

@@ -2,6 +2,7 @@
 from .cluster import manage_cluster
 from .files import manage_files
 from .manager import get_token
+from .region import manage_regions
 from .workspace import get_organization
 from .workspace import get_secret
 from .workspace import get_stage

singlestoredb/tests/test_management.py

@@ -14,6 +14,7 @@
 from singlestoredb.management.job import Status
 from singlestoredb.management.job import TargetType
 from singlestoredb.management.region import Region
+from singlestoredb.management.region import RegionManager
 from singlestoredb.management.utils import NamedList
 
 
@@ -1591,5 +1592,5 @@ def test_no_manager(self):
 
         # Verify from_dict class method
         with self.assertRaises(s2.ManagementError) as cm:
-            Region.get_shared_tier_regions(None)
+            RegionManager.list_shared_tier_regions(None)
         assert 'No workspace manager' in str(cm.exception)