From 61239ad55a2d37435767d3919a09c763f3378774 Mon Sep 17 00:00:00 2001 From: Kevin Smith Date: Thu, 14 May 2026 10:07:34 -0500 Subject: [PATCH 01/14] Fix async UDF event loop starvation under heavy load in Jupyter Async UDFs were running directly in uvicorn's event loop via asyncio.create_task, competing with connection handling under heavy concurrent load. This caused unresponsiveness when running from Jupyter notebooks where the event loop is shared. The fix introduces a dedicated event loop in a background thread for async UDF execution. Coroutines are submitted via run_coroutine_threadsafe() and awaited from the server loop, isolating UDF work from HTTP I/O while preserving cooperative async scheduling between UDFs. Co-Authored-By: Claude Opus 4.6 --- singlestoredb/functions/ext/asgi.py | 38 +++++++++++++++++++++++------ 1 file changed, 30 insertions(+), 8 deletions(-) diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index 331828d1..3929e046 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -1006,6 +1006,15 @@ def __init__( self.log_level = log_level self.disable_metrics = disable_metrics + # Dedicated event loop for async UDF execution, isolated from the server loop + self._udf_loop = asyncio.new_event_loop() + self._udf_thread = threading.Thread( + target=self._udf_loop.run_forever, + daemon=True, + name='async-udf-loop', + ) + self._udf_thread.start() + # Configure logging self._configure_logging() @@ -1039,6 +1048,11 @@ def _configure_logging(self) -> None: # Prevent propagation to avoid duplicate or differently formatted messages self.logger.propagate = False + def shutdown(self) -> None: + """Shut down the dedicated UDF event loop.""" + self._udf_loop.call_soon_threadsafe(self._udf_loop.stop) + self._udf_thread.join(timeout=5) + def get_uvicorn_log_config(self) -> Dict[str, Any]: """ Create uvicorn log config that matches the Application's logging format. @@ -1195,15 +1209,23 @@ async def __call__( func_info['colspec'], b''.join(data), ) - func_task = asyncio.create_task( - func(cancel_event, call_timer, *inputs) - if func_info['is_async'] - else to_thread( - lambda: asyncio.run( - func(cancel_event, call_timer, *inputs), + func_task: 'asyncio.Task[Any]' + if func_info['is_async']: + future = asyncio.run_coroutine_threadsafe( + func(cancel_event, call_timer, *inputs), + self._udf_loop, + ) + func_task = asyncio.create_task( + asyncio.wrap_future(future), # type: ignore[arg-type] + ) + else: + func_task = asyncio.create_task( + to_thread( + lambda: asyncio.run( + func(cancel_event, call_timer, *inputs), + ), ), - ), - ) + ) disconnect_task = asyncio.create_task( asyncio.sleep(int(1e9)) if ignore_cancel else cancel_on_disconnect(receive), From 9fa67f2119fab5f10134addb55d472efe02f7552 Mon Sep 17 00:00:00 2001 From: Kevin Smith Date: Thu, 14 May 2026 10:11:01 -0500 Subject: [PATCH 02/14] Ensure proper cancellation of async UDFs in dedicated loop Cancel the concurrent.futures.Future in the UDF loop on disconnect/timeout so the coroutine is interrupted promptly, not just at the next cancel_on_event row check. Co-Authored-By: Claude Opus 4.6 --- singlestoredb/functions/ext/asgi.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index 3929e046..fa44348d 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -24,6 +24,7 @@ """ import argparse import asyncio +import concurrent.futures import contextvars import dataclasses import datetime @@ -1210,13 +1211,14 @@ async def __call__( ) func_task: 'asyncio.Task[Any]' + udf_future: 'Optional[concurrent.futures.Future[Any]]' = None if func_info['is_async']: - future = asyncio.run_coroutine_threadsafe( + udf_future = asyncio.run_coroutine_threadsafe( func(cancel_event, call_timer, *inputs), self._udf_loop, ) func_task = asyncio.create_task( - asyncio.wrap_future(future), # type: ignore[arg-type] + asyncio.wrap_future(udf_future), # type: ignore[arg-type] ) else: func_task = asyncio.create_task( @@ -1246,12 +1248,16 @@ async def __call__( for task in done: if task is disconnect_task: cancel_event.set() + if udf_future is not None: + udf_future.cancel() raise asyncio.CancelledError( 'Function call was cancelled by client disconnect', ) elif task is timeout_task: cancel_event.set() + if udf_future is not None: + udf_future.cancel() raise asyncio.TimeoutError( 'Function call was cancelled due to timeout', ) From 5ea0b5d7ff12edb8e1aa959544086d2e4bf64a79 Mon Sep 17 00:00:00 2001 From: Kevin Smith Date: Thu, 14 May 2026 10:19:46 -0500 Subject: [PATCH 03/14] Fix create_task expecting coroutine, use ensure_future for wrapped future asyncio.create_task() requires a coroutine but asyncio.wrap_future() returns a Future. Use asyncio.ensure_future() which accepts both. Co-Authored-By: Claude Opus 4.6 --- singlestoredb/functions/ext/asgi.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index fa44348d..ce0c3997 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -1217,8 +1217,8 @@ async def __call__( func(cancel_event, call_timer, *inputs), self._udf_loop, ) - func_task = asyncio.create_task( - asyncio.wrap_future(udf_future), # type: ignore[arg-type] + func_task = asyncio.ensure_future( + asyncio.wrap_future(udf_future), ) else: func_task = asyncio.create_task( From 0a4576c3ad1cc964d42bd65390ed891e2195890d Mon Sep 17 00:00:00 2001 From: Kevin Smith Date: Thu, 14 May 2026 10:43:14 -0500 Subject: [PATCH 04/14] Call Application.shutdown() when replacing UDF server Prevents UDF event loop thread leaks when run_udf_app() is called repeatedly in Jupyter notebooks. Co-Authored-By: Claude Opus 4.6 --- singlestoredb/apps/_python_udfs.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/singlestoredb/apps/_python_udfs.py b/singlestoredb/apps/_python_udfs.py index e45718de..7835307c 100644 --- a/singlestoredb/apps/_python_udfs.py +++ b/singlestoredb/apps/_python_udfs.py @@ -10,8 +10,9 @@ if typing.TYPE_CHECKING: from ._uvicorn_util import AwaitableUvicornServer -# Keep track of currently running server +# Keep track of currently running server and app _running_server: 'typing.Optional[AwaitableUvicornServer]' = None +_running_app: typing.Optional[Application] = None # Maximum number of UDFs allowed MAX_UDFS_LIMIT = 10 @@ -21,7 +22,7 @@ async def run_udf_app( log_level: str = 'error', kill_existing_app_server: bool = True, ) -> UdfConnectionInfo: - global _running_server + global _running_server, _running_app from ._uvicorn_util import AwaitableUvicornServer try: @@ -38,6 +39,9 @@ async def run_udf_app( if _running_server is not None: await _running_server.shutdown() _running_server = None + if _running_app is not None: + _running_app.shutdown() + _running_app = None # Kill if any other process is occupying the port kill_process_by_port(app_config.listen_port) @@ -72,6 +76,7 @@ async def run_udf_app( if app_config.running_interactively: app.register_functions(replace=True) + _running_app = app _running_server = AwaitableUvicornServer(config) asyncio.create_task(_running_server.serve()) await _running_server.wait_for_startup() From 3e1296ed723d3243961475b34b66a1a3aef25ace Mon Sep 17 00:00:00 2001 From: Kevin Smith Date: Thu, 14 May 2026 11:01:30 -0500 Subject: [PATCH 05/14] Add pre-release tag builds with GitHub Release assets Tags matching v*-rc*, v*-test*, v*-alpha*, v*-beta* now trigger the full wheel build pipeline and create a pre-release GitHub Release with all wheels attached. Production releases also attach wheels to the existing release before publishing to PyPI. Co-Authored-By: Claude Opus 4.6 --- .github/workflows/publish.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 5712c3ef..0b64f499 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -207,7 +207,7 @@ jobs: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - name: Publish to PyPI - if: ${{ github.event_name == 'release' || github.event.inputs.publish_pypi == 'true' }} + if: ${{ github.event_name == 'release' || (github.event_name == 'workflow_dispatch' && github.event.inputs.publish_pypi == 'true') }} uses: pypa/gh-action-pypi-publish@release/v1 # - name: Publish Conda package From f557d0e0dad8ebdc297d9e2a0d0ea2ebcd828869 Mon Sep 17 00:00:00 2001 From: Kevin Smith Date: Thu, 14 May 2026 11:23:28 -0500 Subject: [PATCH 06/14] Propagate cancellation to UDF loop and prevent thread leak on failure - Cancel udf_future when func_task is in pending set after asyncio.wait - Cancel udf_future in finally block to ensure cleanup on any exit path - Wrap post-construction code in try/except to call app.shutdown() if validation, config, or registration fails after Application is created Co-Authored-By: Claude Opus 4.6 --- singlestoredb/apps/_python_udfs.py | 44 ++++++++++++++++------------- singlestoredb/functions/ext/asgi.py | 6 ++++ 2 files changed, 30 insertions(+), 20 deletions(-) diff --git a/singlestoredb/apps/_python_udfs.py b/singlestoredb/apps/_python_udfs.py index 7835307c..1039565c 100644 --- a/singlestoredb/apps/_python_udfs.py +++ b/singlestoredb/apps/_python_udfs.py @@ -58,28 +58,32 @@ async def run_udf_app( log_level=log_level, ) - if not app.endpoints: - raise ValueError('You must define at least one function.') - if len(app.endpoints) > MAX_UDFS_LIMIT: - raise ValueError( - f'You can only define a maximum of {MAX_UDFS_LIMIT} functions.', - ) - - config = uvicorn.Config( - app, - host='0.0.0.0', - port=app_config.listen_port, - log_config=app.get_uvicorn_log_config(), - ) + try: + if not app.endpoints: + raise ValueError('You must define at least one function.') + if len(app.endpoints) > MAX_UDFS_LIMIT: + raise ValueError( + f'You can only define a maximum of {MAX_UDFS_LIMIT} functions.', + ) - # Register the functions only if the app is running interactively. - if app_config.running_interactively: - app.register_functions(replace=True) + config = uvicorn.Config( + app, + host='0.0.0.0', + port=app_config.listen_port, + log_config=app.get_uvicorn_log_config(), + ) - _running_app = app - _running_server = AwaitableUvicornServer(config) - asyncio.create_task(_running_server.serve()) - await _running_server.wait_for_startup() + # Register the functions only if the app is running interactively. + if app_config.running_interactively: + app.register_functions(replace=True) + + _running_app = app + _running_server = AwaitableUvicornServer(config) + asyncio.create_task(_running_server.serve()) + await _running_server.wait_for_startup() + except Exception: + app.shutdown() + raise print(f'Python UDF registered at {base_url}') diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index ce0c3997..ce9ab837 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -1244,6 +1244,9 @@ async def __call__( ) await cancel_all_tasks(pending) + if func_task in pending and udf_future is not None: + cancel_event.set() + udf_future.cancel() for task in done: if task is disconnect_task: @@ -1320,6 +1323,9 @@ async def __call__( await send(self.error_response_dict) finally: + if udf_future is not None: + cancel_event.set() + udf_future.cancel() await cancel_all_tasks(all_tasks) # Handle api reflection From c6450f528a15d7d919b24869d7af1e489f989513 Mon Sep 17 00:00:00 2001 From: Kevin Smith Date: Thu, 14 May 2026 12:03:30 -0500 Subject: [PATCH 07/14] Fix udf_future NameError and lazily initialize UDF event loop - Move udf_future initialization before input_handler['load']() to prevent NameError in finally block if parsing raises - Lazily create UDF event loop on first async UDF invocation instead of unconditionally in __init__, avoiding wasted resources for sync-only or metadata-only usage - Guard shutdown() against None loop/thread Co-Authored-By: Claude Opus 4.6 --- singlestoredb/functions/ext/asgi.py | 32 ++++++++++++++++++----------- 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index ce9ab837..2519d948 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -1007,14 +1007,8 @@ def __init__( self.log_level = log_level self.disable_metrics = disable_metrics - # Dedicated event loop for async UDF execution, isolated from the server loop - self._udf_loop = asyncio.new_event_loop() - self._udf_thread = threading.Thread( - target=self._udf_loop.run_forever, - daemon=True, - name='async-udf-loop', - ) - self._udf_thread.start() + self._udf_loop: Optional[asyncio.AbstractEventLoop] = None + self._udf_thread: Optional[threading.Thread] = None # Configure logging self._configure_logging() @@ -1049,10 +1043,24 @@ def _configure_logging(self) -> None: # Prevent propagation to avoid duplicate or differently formatted messages self.logger.propagate = False + def _get_udf_loop(self) -> asyncio.AbstractEventLoop: + """Get or create the dedicated UDF event loop.""" + if self._udf_loop is None: + self._udf_loop = asyncio.new_event_loop() + self._udf_thread = threading.Thread( + target=self._udf_loop.run_forever, + daemon=True, + name='async-udf-loop', + ) + self._udf_thread.start() + return self._udf_loop + def shutdown(self) -> None: """Shut down the dedicated UDF event loop.""" - self._udf_loop.call_soon_threadsafe(self._udf_loop.stop) - self._udf_thread.join(timeout=5) + if self._udf_loop is not None: + self._udf_loop.call_soon_threadsafe(self._udf_loop.stop) + if self._udf_thread is not None: + self._udf_thread.join(timeout=5) def get_uvicorn_log_config(self) -> Dict[str, Any]: """ @@ -1202,6 +1210,7 @@ async def __call__( try: all_tasks = [] result = [] + udf_future: 'Optional[concurrent.futures.Future[Any]]' = None cancel_event = threading.Event() @@ -1211,11 +1220,10 @@ async def __call__( ) func_task: 'asyncio.Task[Any]' - udf_future: 'Optional[concurrent.futures.Future[Any]]' = None if func_info['is_async']: udf_future = asyncio.run_coroutine_threadsafe( func(cancel_event, call_timer, *inputs), - self._udf_loop, + self._get_udf_loop(), ) func_task = asyncio.ensure_future( asyncio.wrap_future(udf_future), From b01445c9df1ab0c285a2b6eb0a280a7b605db240 Mon Sep 17 00:00:00 2001 From: Kevin Smith Date: Thu, 14 May 2026 15:18:49 -0500 Subject: [PATCH 08/14] Reset UDF loop state in shutdown() to allow safe reuse After stopping the event loop and joining the thread, set both _udf_loop and _udf_thread back to None so that _get_udf_loop() can safely recreate them if called after shutdown. Co-Authored-By: Claude Opus 4.6 --- singlestoredb/functions/ext/asgi.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index 2519d948..ba3b17bd 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -1061,6 +1061,8 @@ def shutdown(self) -> None: self._udf_loop.call_soon_threadsafe(self._udf_loop.stop) if self._udf_thread is not None: self._udf_thread.join(timeout=5) + self._udf_loop = None + self._udf_thread = None def get_uvicorn_log_config(self) -> Dict[str, Any]: """ From fadc920b4b4198906bb6f0ea81c7eac42364556a Mon Sep 17 00:00:00 2001 From: Kevin Smith Date: Mon, 18 May 2026 09:50:57 -0500 Subject: [PATCH 09/14] Use thread-per-request for async UDFs instead of shared event loop The dedicated shared event loop still caused starvation under concurrent async UDF calls. Switch to the same model used by sync UDFs: each request gets its own thread with asyncio.run(), eliminating loop contention. Co-Authored-By: Claude Opus 4.6 --- singlestoredb/apps/_python_udfs.py | 51 ++++++++++--------------- singlestoredb/functions/ext/asgi.py | 59 ++++------------------------- 2 files changed, 29 insertions(+), 81 deletions(-) diff --git a/singlestoredb/apps/_python_udfs.py b/singlestoredb/apps/_python_udfs.py index 1039565c..e45718de 100644 --- a/singlestoredb/apps/_python_udfs.py +++ b/singlestoredb/apps/_python_udfs.py @@ -10,9 +10,8 @@ if typing.TYPE_CHECKING: from ._uvicorn_util import AwaitableUvicornServer -# Keep track of currently running server and app +# Keep track of currently running server _running_server: 'typing.Optional[AwaitableUvicornServer]' = None -_running_app: typing.Optional[Application] = None # Maximum number of UDFs allowed MAX_UDFS_LIMIT = 10 @@ -22,7 +21,7 @@ async def run_udf_app( log_level: str = 'error', kill_existing_app_server: bool = True, ) -> UdfConnectionInfo: - global _running_server, _running_app + global _running_server from ._uvicorn_util import AwaitableUvicornServer try: @@ -39,9 +38,6 @@ async def run_udf_app( if _running_server is not None: await _running_server.shutdown() _running_server = None - if _running_app is not None: - _running_app.shutdown() - _running_app = None # Kill if any other process is occupying the port kill_process_by_port(app_config.listen_port) @@ -58,32 +54,27 @@ async def run_udf_app( log_level=log_level, ) - try: - if not app.endpoints: - raise ValueError('You must define at least one function.') - if len(app.endpoints) > MAX_UDFS_LIMIT: - raise ValueError( - f'You can only define a maximum of {MAX_UDFS_LIMIT} functions.', - ) - - config = uvicorn.Config( - app, - host='0.0.0.0', - port=app_config.listen_port, - log_config=app.get_uvicorn_log_config(), + if not app.endpoints: + raise ValueError('You must define at least one function.') + if len(app.endpoints) > MAX_UDFS_LIMIT: + raise ValueError( + f'You can only define a maximum of {MAX_UDFS_LIMIT} functions.', ) - # Register the functions only if the app is running interactively. - if app_config.running_interactively: - app.register_functions(replace=True) - - _running_app = app - _running_server = AwaitableUvicornServer(config) - asyncio.create_task(_running_server.serve()) - await _running_server.wait_for_startup() - except Exception: - app.shutdown() - raise + config = uvicorn.Config( + app, + host='0.0.0.0', + port=app_config.listen_port, + log_config=app.get_uvicorn_log_config(), + ) + + # Register the functions only if the app is running interactively. + if app_config.running_interactively: + app.register_functions(replace=True) + + _running_server = AwaitableUvicornServer(config) + asyncio.create_task(_running_server.serve()) + await _running_server.wait_for_startup() print(f'Python UDF registered at {base_url}') diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index ba3b17bd..b22f1c26 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -24,7 +24,6 @@ """ import argparse import asyncio -import concurrent.futures import contextvars import dataclasses import datetime @@ -1007,9 +1006,6 @@ def __init__( self.log_level = log_level self.disable_metrics = disable_metrics - self._udf_loop: Optional[asyncio.AbstractEventLoop] = None - self._udf_thread: Optional[threading.Thread] = None - # Configure logging self._configure_logging() @@ -1043,27 +1039,6 @@ def _configure_logging(self) -> None: # Prevent propagation to avoid duplicate or differently formatted messages self.logger.propagate = False - def _get_udf_loop(self) -> asyncio.AbstractEventLoop: - """Get or create the dedicated UDF event loop.""" - if self._udf_loop is None: - self._udf_loop = asyncio.new_event_loop() - self._udf_thread = threading.Thread( - target=self._udf_loop.run_forever, - daemon=True, - name='async-udf-loop', - ) - self._udf_thread.start() - return self._udf_loop - - def shutdown(self) -> None: - """Shut down the dedicated UDF event loop.""" - if self._udf_loop is not None: - self._udf_loop.call_soon_threadsafe(self._udf_loop.stop) - if self._udf_thread is not None: - self._udf_thread.join(timeout=5) - self._udf_loop = None - self._udf_thread = None - def get_uvicorn_log_config(self) -> Dict[str, Any]: """ Create uvicorn log config that matches the Application's logging format. @@ -1212,7 +1187,6 @@ async def __call__( try: all_tasks = [] result = [] - udf_future: 'Optional[concurrent.futures.Future[Any]]' = None cancel_event = threading.Event() @@ -1221,23 +1195,13 @@ async def __call__( func_info['colspec'], b''.join(data), ) - func_task: 'asyncio.Task[Any]' - if func_info['is_async']: - udf_future = asyncio.run_coroutine_threadsafe( - func(cancel_event, call_timer, *inputs), - self._get_udf_loop(), - ) - func_task = asyncio.ensure_future( - asyncio.wrap_future(udf_future), - ) - else: - func_task = asyncio.create_task( - to_thread( - lambda: asyncio.run( - func(cancel_event, call_timer, *inputs), - ), + func_task = asyncio.create_task( + to_thread( + lambda: asyncio.run( + func(cancel_event, call_timer, *inputs), ), - ) + ), + ) disconnect_task = asyncio.create_task( asyncio.sleep(int(1e9)) if ignore_cancel else cancel_on_disconnect(receive), @@ -1254,23 +1218,18 @@ async def __call__( ) await cancel_all_tasks(pending) - if func_task in pending and udf_future is not None: + if func_task in pending: cancel_event.set() - udf_future.cancel() for task in done: if task is disconnect_task: cancel_event.set() - if udf_future is not None: - udf_future.cancel() raise asyncio.CancelledError( 'Function call was cancelled by client disconnect', ) elif task is timeout_task: cancel_event.set() - if udf_future is not None: - udf_future.cancel() raise asyncio.TimeoutError( 'Function call was cancelled due to timeout', ) @@ -1333,9 +1292,7 @@ async def __call__( await send(self.error_response_dict) finally: - if udf_future is not None: - cancel_event.set() - udf_future.cancel() + cancel_event.set() await cancel_all_tasks(all_tasks) # Handle api reflection From 1ef852ee4995f0037f98cbba6eb48b3b36001907 Mon Sep 17 00:00:00 2001 From: Kevin Smith Date: Mon, 18 May 2026 10:12:09 -0500 Subject: [PATCH 10/14] Add cancellable wrapper for responsive async UDF cancellation Wraps the inner coroutine in _cancellable_run which polls cancel_event and raises CancelledError at the next await (~100ms), ensuring vector UDFs respect disconnect/timeout signals without waiting for completion. Co-Authored-By: Claude Opus 4.6 --- singlestoredb/functions/ext/asgi.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index b22f1c26..e426a800 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -113,6 +113,28 @@ async def to_thread( return await loop.run_in_executor(None, func_call) +async def _poll_cancel(cancel_event: threading.Event) -> None: + while not cancel_event.is_set(): + await asyncio.sleep(0.1) + + +async def _cancellable_run( + cancel_event: threading.Event, + coro: Any, +) -> Any: + task = asyncio.create_task(coro) + cancel_check = asyncio.create_task(_poll_cancel(cancel_event)) + done, pending = await asyncio.wait( + [task, cancel_check], return_when=asyncio.FIRST_COMPLETED, + ) + for p in pending: + p.cancel() + if cancel_check in done: + task.cancel() + raise asyncio.CancelledError() + return task.result() + + # Use negative values to indicate unsigned ints / binary data / usec time precision rowdat_1_type_map = { 'bool': ft.LONGLONG, @@ -1198,7 +1220,10 @@ async def __call__( func_task = asyncio.create_task( to_thread( lambda: asyncio.run( - func(cancel_event, call_timer, *inputs), + _cancellable_run( + cancel_event, + func(cancel_event, call_timer, *inputs), + ), ), ), ) From 4c2caff3736b9b168abc282244cc706a5e5c01c9 Mon Sep 17 00:00:00 2001 From: Kevin Smith Date: Tue, 19 May 2026 09:00:21 -0500 Subject: [PATCH 11/14] Fix event loop closed error and add comprehensive UDF dispatch tests Replace asyncio.run() with _run_with_graceful_shutdown() that drains pending callbacks before closing the loop, preventing RuntimeError from httpx/anyio TLS cleanup in async UDFs calling OpenAI/LangChain APIs. Add 17 unit tests covering graceful shutdown, cancellation timing, exception propagation, context variable isolation, and concurrent safety. Co-Authored-By: Claude Opus 4.6 --- singlestoredb/functions/ext/asgi.py | 30 ++- singlestoredb/tests/test_udf_event_loop.py | 296 +++++++++++++++++++++ 2 files changed, 325 insertions(+), 1 deletion(-) create mode 100644 singlestoredb/tests/test_udf_event_loop.py diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index e426a800..329cae64 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -135,6 +135,34 @@ async def _cancellable_run( return task.result() +def _run_with_graceful_shutdown(coro: Any) -> Any: + """Run a coroutine in a new event loop, draining callbacks before close. + + Unlike asyncio.run(), this prevents 'Event loop is closed' errors from + libraries (httpx/anyio) that schedule cleanup callbacks during teardown. + """ + loop = asyncio.new_event_loop() + try: + asyncio.set_event_loop(loop) + return loop.run_until_complete(coro) + finally: + try: + pending = asyncio.all_tasks(loop) + if pending: + for task in pending: + task.cancel() + loop.run_until_complete( + asyncio.gather(*pending, return_exceptions=True), + ) + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.run_until_complete(loop.shutdown_default_executor()) + finally: + loop.call_soon(loop.stop) + loop.run_forever() + asyncio.set_event_loop(None) + loop.close() + + # Use negative values to indicate unsigned ints / binary data / usec time precision rowdat_1_type_map = { 'bool': ft.LONGLONG, @@ -1219,7 +1247,7 @@ async def __call__( func_task = asyncio.create_task( to_thread( - lambda: asyncio.run( + lambda: _run_with_graceful_shutdown( _cancellable_run( cancel_event, func(cancel_event, call_timer, *inputs), diff --git a/singlestoredb/tests/test_udf_event_loop.py b/singlestoredb/tests/test_udf_event_loop.py new file mode 100644 index 00000000..279f4a18 --- /dev/null +++ b/singlestoredb/tests/test_udf_event_loop.py @@ -0,0 +1,296 @@ +"""Tests for async UDF event loop graceful shutdown.""" +import asyncio +import contextvars +import threading +import time +import unittest +from typing import Any +from typing import List + +from ..functions.ext.asgi import _cancellable_run +from ..functions.ext.asgi import _run_with_graceful_shutdown +from ..functions.ext.asgi import to_thread + + +class TestRunWithGracefulShutdown(unittest.TestCase): + """Test _run_with_graceful_shutdown handles loop cleanup properly.""" + + def test_basic_coroutine(self) -> None: + async def simple() -> int: + return 42 + + result = _run_with_graceful_shutdown(simple()) + self.assertEqual(result, 42) + + def test_callbacks_drained_before_close(self) -> None: + """Simulate httpx/anyio scheduling call_soon during teardown. + + This is the exact pattern that causes 'Event loop is closed' with + asyncio.run() -- a library schedules a callback in its __del__ or + aclose() that fires after the loop is closed. + """ + callback_executed: List[bool] = [] + + async def coroutine_with_cleanup_callback() -> str: + loop = asyncio.get_running_loop() + loop.call_soon(lambda: callback_executed.append(True)) + return 'done' + + result = _run_with_graceful_shutdown(coroutine_with_cleanup_callback()) + self.assertEqual(result, 'done') + self.assertEqual(callback_executed, [True]) + + def test_no_event_loop_closed_error(self) -> None: + """Verify no RuntimeError when cleanup schedules on the loop.""" + errors: List[RuntimeError] = [] + + async def simulate_httpx_teardown() -> str: + loop = asyncio.get_running_loop() + + def deferred_cleanup() -> None: + try: + loop.call_soon(lambda: None) + except RuntimeError as e: + errors.append(e) + + loop.call_soon(deferred_cleanup) + return 'ok' + + result = _run_with_graceful_shutdown(simulate_httpx_teardown()) + self.assertEqual(result, 'ok') + self.assertEqual(errors, []) + + def test_exception_propagates(self) -> None: + async def failing() -> None: + raise ValueError('test error') + + with self.assertRaises(ValueError) as ctx: + _run_with_graceful_shutdown(failing()) + self.assertEqual(str(ctx.exception), 'test error') + + def test_callbacks_drained_even_on_exception(self) -> None: + """Cleanup callbacks still run even if coroutine raises.""" + callback_executed: List[bool] = [] + + async def failing_with_callback() -> None: + loop = asyncio.get_running_loop() + loop.call_soon(lambda: callback_executed.append(True)) + raise ValueError('boom') + + with self.assertRaises(ValueError): + _run_with_graceful_shutdown(failing_with_callback()) + self.assertEqual(callback_executed, [True]) + + def test_pending_tasks_cancelled(self) -> None: + """Background tasks are cancelled during shutdown.""" + async def background() -> None: + await asyncio.sleep(999) + + async def main_with_background_task() -> str: + asyncio.create_task(background()) + return 'done' + + result = _run_with_graceful_shutdown(main_with_background_task()) + self.assertEqual(result, 'done') + + def test_isolation_between_calls(self) -> None: + """Each call gets its own event loop that is closed after use.""" + loops: List[asyncio.AbstractEventLoop] = [] + + async def capture_loop() -> bool: + loops.append(asyncio.get_running_loop()) + return True + + _run_with_graceful_shutdown(capture_loop()) + first_loop = loops[0] + self.assertTrue(first_loop.is_closed()) + + _run_with_graceful_shutdown(capture_loop()) + second_loop = loops[1] + self.assertTrue(second_loop.is_closed()) + + def test_cancellable_run_integration(self) -> None: + """Verify _cancellable_run works inside _run_with_graceful_shutdown.""" + cancel_event = threading.Event() + + async def slow_func() -> str: + return 'completed' + + result = _run_with_graceful_shutdown( + _cancellable_run(cancel_event, slow_func()), + ) + self.assertEqual(result, 'completed') + + def test_cancellation_via_event(self) -> None: + """Verify cancellation propagates through the full stack.""" + cancel_event = threading.Event() + cancel_event.set() + + async def blocked_func() -> str: + await asyncio.sleep(999) + return 'should not reach' + + with self.assertRaises(asyncio.CancelledError): + _run_with_graceful_shutdown( + _cancellable_run(cancel_event, blocked_func()), + ) + + +class TestUDFDispatchEdgeCases(unittest.TestCase): + """Test edge cases in the UDF dispatch stack.""" + + def test_timeout_cancels_running_function(self) -> None: + """Cancel event set from timer thread cancels a blocked coroutine.""" + cancel_event = threading.Event() + + async def long_running() -> str: + await asyncio.sleep(999) + return 'should not reach' + + def set_cancel_after_delay() -> None: + time.sleep(0.2) + cancel_event.set() + + timer = threading.Thread(target=set_cancel_after_delay) + timer.start() + + start = time.monotonic() + with self.assertRaises(asyncio.CancelledError): + _run_with_graceful_shutdown( + _cancellable_run(cancel_event, long_running()), + ) + elapsed = time.monotonic() - start + timer.join() + # 0.2s delay + up to 0.1s poll interval + margin + self.assertLess(elapsed, 0.5) + + def test_exception_propagates_through_full_stack(self) -> None: + """User exception propagates unwrapped through the entire dispatch.""" + cancel_event = threading.Event() + + class CustomUDFError(Exception): + pass + + async def failing_udf() -> None: + raise CustomUDFError('embedding service unavailable') + + with self.assertRaises(CustomUDFError) as ctx: + _run_with_graceful_shutdown( + _cancellable_run(cancel_event, failing_udf()), + ) + self.assertEqual(str(ctx.exception), 'embedding service unavailable') + + def test_cancel_event_detected_within_poll_interval(self) -> None: + """Cancellation is detected within one poll cycle (0.1s).""" + cancel_event = threading.Event() + + async def blocked() -> str: + await asyncio.sleep(999) + return 'unreachable' + + def set_cancel() -> None: + time.sleep(0.05) + cancel_event.set() + + timer = threading.Thread(target=set_cancel) + timer.start() + + start = time.monotonic() + with self.assertRaises(asyncio.CancelledError): + _run_with_graceful_shutdown( + _cancellable_run(cancel_event, blocked()), + ) + elapsed = time.monotonic() - start + timer.join() + # 0.05s delay + 0.1s poll interval + margin + self.assertLess(elapsed, 0.25) + + def test_context_vars_propagate_through_to_thread(self) -> None: + """Context variables are visible inside to_thread executor.""" + test_var: contextvars.ContextVar[str] = contextvars.ContextVar( + 'test_var', + ) + test_var.set('hello_from_parent') + captured: List[str] = [] + + def read_context_var() -> str: + val = test_var.get('NOT_FOUND') + captured.append(val) + return val + + async def run_in_thread() -> str: + return await to_thread(read_context_var) + + result = _run_with_graceful_shutdown(run_in_thread()) + self.assertEqual(result, 'hello_from_parent') + self.assertEqual(captured, ['hello_from_parent']) + + def test_concurrent_requests_isolated(self) -> None: + """Parallel executions don't share state.""" + results: List[Any] = [None, None, None] + + def run_isolated(index: int) -> None: + async def compute() -> int: + await asyncio.sleep(0.05) + return index * 10 + + results[index] = _run_with_graceful_shutdown(compute()) + + threads = [ + threading.Thread(target=run_isolated, args=(i,)) + for i in range(3) + ] + for t in threads: + t.start() + for t in threads: + t.join() + + self.assertEqual(results, [0, 10, 20]) + + def test_sync_function_through_async_wrapper(self) -> None: + """Synchronous function works when wrapped as async coroutine.""" + cancel_event = threading.Event() + + async def sync_as_async() -> int: + # Simulates what decorator.py's async_wrapper does for sync UDFs + return 42 + 1 + + result = _run_with_graceful_shutdown( + _cancellable_run(cancel_event, sync_as_async()), + ) + self.assertEqual(result, 43) + + def test_cancel_event_not_set_on_success(self) -> None: + """Cancel event remains unset after successful execution.""" + cancel_event = threading.Event() + + async def quick() -> str: + return 'fast' + + result = _run_with_graceful_shutdown( + _cancellable_run(cancel_event, quick()), + ) + self.assertEqual(result, 'fast') + self.assertFalse(cancel_event.is_set()) + + def test_callbacks_from_cancelled_tasks_still_drain(self) -> None: + """Background task callbacks drain even when task is cancelled.""" + drained: List[bool] = [] + + async def bg_with_callback() -> None: + loop = asyncio.get_running_loop() + loop.call_soon(lambda: drained.append(True)) + await asyncio.sleep(999) + + async def main() -> str: + asyncio.create_task(bg_with_callback()) + await asyncio.sleep(0.05) # Let background task start + return 'done' + + result = _run_with_graceful_shutdown(main()) + self.assertEqual(result, 'done') + self.assertEqual(drained, [True]) + + +if __name__ == '__main__': + unittest.main() From 0070633dece676d92ee80bd0ecd51bf67f15303c Mon Sep 17 00:00:00 2001 From: Kevin Smith Date: Thu, 28 May 2026 08:42:24 -0500 Subject: [PATCH 12/14] Add PEP 440 pre-release version patching to publish workflow Parses git tag suffixes (-test, -alpha, -beta, -rc) and patches pyproject.toml with the corresponding PEP 440 version before building wheels. Full releases (no suffix) are unaffected. Co-Authored-By: Claude Opus 4.6 --- .github/workflows/publish.yml | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 0b64f499..c4e92ef8 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -75,6 +75,33 @@ jobs: steps: - uses: actions/checkout@v3 + - name: Patch version for pre-release tag + if: github.event_name == 'push' + shell: python + run: | + import re, os, pathlib + tag = os.environ["GITHUB_REF_NAME"].lstrip("v") + patterns = [ + (r"^(\d+\.\d+\.\d+)-test(\d+)$", r"\1.dev\2"), + (r"^(\d+\.\d+\.\d+)-alpha(\d+)$", r"\1a\2"), + (r"^(\d+\.\d+\.\d+)-beta(\d+)$", r"\1b\2"), + (r"^(\d+\.\d+\.\d+)-rc(\d+)$", r"\1rc\2"), + ] + version = None + for pattern, repl in patterns: + m = re.match(pattern, tag) + if m: + version = re.sub(pattern, repl, tag) + break + if version is None: + print("No pre-release suffix, keeping version as-is") + else: + print(f"Patching version to: {version}") + p = pathlib.Path("pyproject.toml") + content = p.read_text() + content = re.sub(r'^version = ".*"', f'version = "{version}"', content, count=1, flags=re.MULTILINE) + p.write_text(content) + - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v4 with: From 78a96f844ecd569454fa5a2ff5d3d8110d725c35 Mon Sep 17 00:00:00 2001 From: Ankit Saini Date: Wed, 3 Jun 2026 23:50:15 +0530 Subject: [PATCH 13/14] reuse event loop across requests on same thread Co-authored-by: Cursor --- singlestoredb/functions/ext/asgi.py | 86 +++++-- singlestoredb/tests/test_udf_event_loop.py | 286 ++++++++++----------- 2 files changed, 203 insertions(+), 169 deletions(-) diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index 329cae64..294b17c2 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -24,6 +24,7 @@ """ import argparse import asyncio +import atexit import contextvars import dataclasses import datetime @@ -135,32 +136,67 @@ async def _cancellable_run( return task.result() -def _run_with_graceful_shutdown(coro: Any) -> Any: - """Run a coroutine in a new event loop, draining callbacks before close. +# Each `to_thread` worker thread owns a long-lived event loop reused across +# requests, so loop-bound resources (HTTP pools, DB sessions, sockets) can +# survive between calls handled by the same thread. +_thread_local = threading.local() +_loop_registry: 'Set[asyncio.AbstractEventLoop]' = set() +_loop_registry_lock = threading.Lock() - Unlike asyncio.run(), this prevents 'Event loop is closed' errors from - libraries (httpx/anyio) that schedule cleanup callbacks during teardown. - """ - loop = asyncio.new_event_loop() - try: + +def _get_thread_loop() -> asyncio.AbstractEventLoop: + """Return (creating if needed) the calling thread's persistent loop.""" + loop = getattr(_thread_local, 'loop', None) + if loop is None or loop.is_closed(): + loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) - return loop.run_until_complete(coro) - finally: + _thread_local.loop = loop + with _loop_registry_lock: + _loop_registry.add(loop) + return loop + + +def _run_on_thread_loop(coro: Any) -> Any: + """ + Run ``coro`` on the calling thread's persistent loop. + + The loop is never closed between calls, so loop-bound resources (e.g. + httpx keep-alive pools) survive across requests and the deferred + "Event loop is closed" errors thrown by httpx/anyio at teardown do not + occur. + + Caveat: tasks the user code spawns via ``asyncio.create_task`` and + leaves running outlive the current call too. That is the price of + keeping shared resources alive; ``cancel_event`` does not reach them. + """ + loop = _get_thread_loop() + return loop.run_until_complete(coro) + + +def _shutdown_thread_loops() -> None: + """Best-effort cleanup of all persistent worker-thread loops at exit.""" + with _loop_registry_lock: + loops = list(_loop_registry) + _loop_registry.clear() + + for loop in loops: + if loop.is_closed(): + continue try: - pending = asyncio.all_tasks(loop) - if pending: - for task in pending: - task.cancel() - loop.run_until_complete( - asyncio.gather(*pending, return_exceptions=True), - ) + # Owning thread is no longer running the loop; safe to drive + # teardown from this (exiting) thread. loop.run_until_complete(loop.shutdown_asyncgens()) loop.run_until_complete(loop.shutdown_default_executor()) + except Exception: + pass finally: - loop.call_soon(loop.stop) - loop.run_forever() - asyncio.set_event_loop(None) - loop.close() + try: + loop.close() + except Exception: + pass + + +atexit.register(_shutdown_thread_loops) # Use negative values to indicate unsigned ints / binary data / usec time precision @@ -1247,7 +1283,7 @@ async def __call__( func_task = asyncio.create_task( to_thread( - lambda: _run_with_graceful_shutdown( + lambda: _run_on_thread_loop( _cancellable_run( cancel_event, func(cancel_event, call_timer, *inputs), @@ -1270,19 +1306,21 @@ async def __call__( all_tasks, return_when=asyncio.FIRST_COMPLETED, ) - await cancel_all_tasks(pending) + # Signal the worker before awaiting cancellation: cancelling + # func_task only flips its asyncio wrapper, not the executor + # work; only cancel_event reaches the worker loop. if func_task in pending: cancel_event.set() + await cancel_all_tasks(pending) + for task in done: if task is disconnect_task: - cancel_event.set() raise asyncio.CancelledError( 'Function call was cancelled by client disconnect', ) elif task is timeout_task: - cancel_event.set() raise asyncio.TimeoutError( 'Function call was cancelled due to timeout', ) diff --git a/singlestoredb/tests/test_udf_event_loop.py b/singlestoredb/tests/test_udf_event_loop.py index 279f4a18..ace45211 100644 --- a/singlestoredb/tests/test_udf_event_loop.py +++ b/singlestoredb/tests/test_udf_event_loop.py @@ -1,4 +1,4 @@ -"""Tests for async UDF event loop graceful shutdown.""" +"""Tests for the async UDF persistent per-thread event loop.""" import asyncio import contextvars import threading @@ -8,134 +8,11 @@ from typing import List from ..functions.ext.asgi import _cancellable_run -from ..functions.ext.asgi import _run_with_graceful_shutdown +from ..functions.ext.asgi import _get_thread_loop +from ..functions.ext.asgi import _run_on_thread_loop from ..functions.ext.asgi import to_thread -class TestRunWithGracefulShutdown(unittest.TestCase): - """Test _run_with_graceful_shutdown handles loop cleanup properly.""" - - def test_basic_coroutine(self) -> None: - async def simple() -> int: - return 42 - - result = _run_with_graceful_shutdown(simple()) - self.assertEqual(result, 42) - - def test_callbacks_drained_before_close(self) -> None: - """Simulate httpx/anyio scheduling call_soon during teardown. - - This is the exact pattern that causes 'Event loop is closed' with - asyncio.run() -- a library schedules a callback in its __del__ or - aclose() that fires after the loop is closed. - """ - callback_executed: List[bool] = [] - - async def coroutine_with_cleanup_callback() -> str: - loop = asyncio.get_running_loop() - loop.call_soon(lambda: callback_executed.append(True)) - return 'done' - - result = _run_with_graceful_shutdown(coroutine_with_cleanup_callback()) - self.assertEqual(result, 'done') - self.assertEqual(callback_executed, [True]) - - def test_no_event_loop_closed_error(self) -> None: - """Verify no RuntimeError when cleanup schedules on the loop.""" - errors: List[RuntimeError] = [] - - async def simulate_httpx_teardown() -> str: - loop = asyncio.get_running_loop() - - def deferred_cleanup() -> None: - try: - loop.call_soon(lambda: None) - except RuntimeError as e: - errors.append(e) - - loop.call_soon(deferred_cleanup) - return 'ok' - - result = _run_with_graceful_shutdown(simulate_httpx_teardown()) - self.assertEqual(result, 'ok') - self.assertEqual(errors, []) - - def test_exception_propagates(self) -> None: - async def failing() -> None: - raise ValueError('test error') - - with self.assertRaises(ValueError) as ctx: - _run_with_graceful_shutdown(failing()) - self.assertEqual(str(ctx.exception), 'test error') - - def test_callbacks_drained_even_on_exception(self) -> None: - """Cleanup callbacks still run even if coroutine raises.""" - callback_executed: List[bool] = [] - - async def failing_with_callback() -> None: - loop = asyncio.get_running_loop() - loop.call_soon(lambda: callback_executed.append(True)) - raise ValueError('boom') - - with self.assertRaises(ValueError): - _run_with_graceful_shutdown(failing_with_callback()) - self.assertEqual(callback_executed, [True]) - - def test_pending_tasks_cancelled(self) -> None: - """Background tasks are cancelled during shutdown.""" - async def background() -> None: - await asyncio.sleep(999) - - async def main_with_background_task() -> str: - asyncio.create_task(background()) - return 'done' - - result = _run_with_graceful_shutdown(main_with_background_task()) - self.assertEqual(result, 'done') - - def test_isolation_between_calls(self) -> None: - """Each call gets its own event loop that is closed after use.""" - loops: List[asyncio.AbstractEventLoop] = [] - - async def capture_loop() -> bool: - loops.append(asyncio.get_running_loop()) - return True - - _run_with_graceful_shutdown(capture_loop()) - first_loop = loops[0] - self.assertTrue(first_loop.is_closed()) - - _run_with_graceful_shutdown(capture_loop()) - second_loop = loops[1] - self.assertTrue(second_loop.is_closed()) - - def test_cancellable_run_integration(self) -> None: - """Verify _cancellable_run works inside _run_with_graceful_shutdown.""" - cancel_event = threading.Event() - - async def slow_func() -> str: - return 'completed' - - result = _run_with_graceful_shutdown( - _cancellable_run(cancel_event, slow_func()), - ) - self.assertEqual(result, 'completed') - - def test_cancellation_via_event(self) -> None: - """Verify cancellation propagates through the full stack.""" - cancel_event = threading.Event() - cancel_event.set() - - async def blocked_func() -> str: - await asyncio.sleep(999) - return 'should not reach' - - with self.assertRaises(asyncio.CancelledError): - _run_with_graceful_shutdown( - _cancellable_run(cancel_event, blocked_func()), - ) - - class TestUDFDispatchEdgeCases(unittest.TestCase): """Test edge cases in the UDF dispatch stack.""" @@ -156,7 +33,7 @@ def set_cancel_after_delay() -> None: start = time.monotonic() with self.assertRaises(asyncio.CancelledError): - _run_with_graceful_shutdown( + _run_on_thread_loop( _cancellable_run(cancel_event, long_running()), ) elapsed = time.monotonic() - start @@ -175,7 +52,7 @@ async def failing_udf() -> None: raise CustomUDFError('embedding service unavailable') with self.assertRaises(CustomUDFError) as ctx: - _run_with_graceful_shutdown( + _run_on_thread_loop( _cancellable_run(cancel_event, failing_udf()), ) self.assertEqual(str(ctx.exception), 'embedding service unavailable') @@ -197,7 +74,7 @@ def set_cancel() -> None: start = time.monotonic() with self.assertRaises(asyncio.CancelledError): - _run_with_graceful_shutdown( + _run_on_thread_loop( _cancellable_run(cancel_event, blocked()), ) elapsed = time.monotonic() - start @@ -221,7 +98,7 @@ def read_context_var() -> str: async def run_in_thread() -> str: return await to_thread(read_context_var) - result = _run_with_graceful_shutdown(run_in_thread()) + result = _run_on_thread_loop(run_in_thread()) self.assertEqual(result, 'hello_from_parent') self.assertEqual(captured, ['hello_from_parent']) @@ -234,7 +111,7 @@ async def compute() -> int: await asyncio.sleep(0.05) return index * 10 - results[index] = _run_with_graceful_shutdown(compute()) + results[index] = _run_on_thread_loop(compute()) threads = [ threading.Thread(target=run_isolated, args=(i,)) @@ -255,7 +132,7 @@ async def sync_as_async() -> int: # Simulates what decorator.py's async_wrapper does for sync UDFs return 42 + 1 - result = _run_with_graceful_shutdown( + result = _run_on_thread_loop( _cancellable_run(cancel_event, sync_as_async()), ) self.assertEqual(result, 43) @@ -267,29 +144,148 @@ def test_cancel_event_not_set_on_success(self) -> None: async def quick() -> str: return 'fast' - result = _run_with_graceful_shutdown( + result = _run_on_thread_loop( _cancellable_run(cancel_event, quick()), ) self.assertEqual(result, 'fast') self.assertFalse(cancel_event.is_set()) - def test_callbacks_from_cancelled_tasks_still_drain(self) -> None: - """Background task callbacks drain even when task is cancelled.""" - drained: List[bool] = [] - async def bg_with_callback() -> None: +class TestRunOnThreadLoop(unittest.TestCase): + """Test _run_on_thread_loop reuses a persistent per-thread event loop.""" + + def test_basic_coroutine(self) -> None: + async def simple() -> int: + return 42 + + self.assertEqual(_run_on_thread_loop(simple()), 42) + + def test_loop_reused_across_calls(self) -> None: + """The same loop object is reused for successive calls in a thread.""" + loops: List[asyncio.AbstractEventLoop] = [] + + async def capture_loop() -> bool: + loops.append(asyncio.get_running_loop()) + return True + + _run_on_thread_loop(capture_loop()) + _run_on_thread_loop(capture_loop()) + + self.assertIs(loops[0], loops[1]) + + def test_loop_not_closed_between_calls(self) -> None: + """The persistent loop stays open so resources survive requests.""" + captured: List[asyncio.AbstractEventLoop] = [] + + async def capture_loop() -> bool: + captured.append(asyncio.get_running_loop()) + return True + + _run_on_thread_loop(capture_loop()) + loop = captured[0] + self.assertFalse(loop.is_closed()) + + # Still usable for the next request. + _run_on_thread_loop(capture_loop()) + self.assertFalse(loop.is_closed()) + + def test_async_resource_survives_between_calls(self) -> None: + """An object bound to the loop can be reused on the next call. + + This mirrors caching e.g. an httpx.AsyncClient keyed by the loop and + reusing its connection pool on subsequent requests. + """ + clients: dict = {} + + async def get_or_create_client() -> int: loop = asyncio.get_running_loop() - loop.call_soon(lambda: drained.append(True)) + if loop not in clients: + clients[loop] = object() + return id(clients[loop]) + + first = _run_on_thread_loop(get_or_create_client()) + second = _run_on_thread_loop(get_or_create_client()) + + self.assertEqual(first, second) + self.assertEqual(len(clients), 1) + + def test_separate_threads_get_separate_loops(self) -> None: + """Each worker thread owns its own persistent loop.""" + loops: List[asyncio.AbstractEventLoop] = [] + lock = threading.Lock() + + def run_in_thread() -> None: + async def capture() -> bool: + with lock: + loops.append(asyncio.get_running_loop()) + return True + + _run_on_thread_loop(capture()) + + threads = [threading.Thread(target=run_in_thread) for _ in range(3)] + for t in threads: + t.start() + for t in threads: + t.join() + + self.assertEqual(len(loops), 3) + self.assertEqual(len({id(loop) for loop in loops}), 3) + + def test_get_thread_loop_idempotent(self) -> None: + """_get_thread_loop returns the same loop on repeated calls.""" + def run_in_thread(out: List[asyncio.AbstractEventLoop]) -> None: + out.append(_get_thread_loop()) + out.append(_get_thread_loop()) + + out: List[asyncio.AbstractEventLoop] = [] + t = threading.Thread(target=run_in_thread, args=(out,)) + t.start() + t.join() + + self.assertIs(out[0], out[1]) + + def test_exception_propagates(self) -> None: + async def failing() -> None: + raise ValueError('test error') + + with self.assertRaises(ValueError) as ctx: + _run_on_thread_loop(failing()) + self.assertEqual(str(ctx.exception), 'test error') + + def test_cancellable_run_integration(self) -> None: + """_cancellable_run works on the persistent loop.""" + cancel_event = threading.Event() + + async def slow_func() -> str: + return 'completed' + + result = _run_on_thread_loop( + _cancellable_run(cancel_event, slow_func()), + ) + self.assertEqual(result, 'completed') + + def test_cancellation_via_event(self) -> None: + """Cancellation propagates through the persistent-loop stack.""" + cancel_event = threading.Event() + cancel_event.set() + + async def blocked_func() -> str: await asyncio.sleep(999) + return 'should not reach' + + with self.assertRaises(asyncio.CancelledError): + _run_on_thread_loop( + _cancellable_run(cancel_event, blocked_func()), + ) - async def main() -> str: - asyncio.create_task(bg_with_callback()) - await asyncio.sleep(0.05) # Let background task start - return 'done' + # Loop must remain usable after a cancelled request. + async def quick() -> str: + return 'ok' - result = _run_with_graceful_shutdown(main()) - self.assertEqual(result, 'done') - self.assertEqual(drained, [True]) + self.assertEqual( + _run_on_thread_loop(_cancellable_run(threading.Event(), quick())), + 'ok', + ) if __name__ == '__main__': From 4c8958b76bbe3ae27ca8e257185e57b013633a00 Mon Sep 17 00:00:00 2001 From: Kaushik Kampli Date: Thu, 4 Jun 2026 10:12:39 +0530 Subject: [PATCH 14/14] lint --- singlestoredb/tests/test_udf_event_loop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/singlestoredb/tests/test_udf_event_loop.py b/singlestoredb/tests/test_udf_event_loop.py index ace45211..bde47389 100644 --- a/singlestoredb/tests/test_udf_event_loop.py +++ b/singlestoredb/tests/test_udf_event_loop.py @@ -195,7 +195,7 @@ def test_async_resource_survives_between_calls(self) -> None: This mirrors caching e.g. an httpx.AsyncClient keyed by the loop and reusing its connection pool on subsequent requests. """ - clients: dict = {} + clients: dict[asyncio.AbstractEventLoop, object] = {} async def get_or_create_client() -> int: loop = asyncio.get_running_loop()