Run each async udf in its own thread

.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:
@@ -207,7 +234,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

singlestoredb/functions/ext/asgi.py

@@ -24,6 +24,7 @@
 """
 import argparse
 import asyncio
+import atexit
 import contextvars
 import dataclasses
 import datetime
@@ -113,6 +114,91 @@ 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()
+
+
+# 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()
+
+
+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)
+        _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:
+            # 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:
+            try:
+                loop.close()
+            except Exception:
+                pass
+
+
+atexit.register(_shutdown_thread_loops)
+
+
 # Use negative values to indicate unsigned ints / binary data / usec time precision
 rowdat_1_type_map = {
     'bool': ft.LONGLONG,
@@ -1196,11 +1282,12 @@ async def __call__(
                     )
 
                 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),
+                    to_thread(
+                        lambda: _run_on_thread_loop(
+                            _cancellable_run(
+                                cancel_event,
+                                func(cancel_event, call_timer, *inputs),
+                            ),
                         ),
                     ),
                 )
@@ -1219,17 +1306,21 @@ async def __call__(
                         all_tasks, return_when=asyncio.FIRST_COMPLETED,
                     )
 
+                # 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',
                         )
@@ -1292,6 +1383,7 @@ async def __call__(
                 await send(self.error_response_dict)
 
             finally:
+                cancel_event.set()
                 await cancel_all_tasks(all_tasks)
 
         # Handle api reflection