Use to_thread for sync transaction functions (#454)

Otherwise, the sync transaction functions will block the entire event
loop.

Author: qianl15

dbos/_core.py

@@ -56,6 +56,7 @@
     DBOSWorkflowConflictIDError,
     DBOSWorkflowFunctionNotFoundError,
 )
+from ._logger import dbos_logger
 from ._registrations import (
     DEFAULT_MAX_RECOVERY_ATTEMPTS,
     get_config_name,
@@ -96,6 +97,14 @@
 TEMP_SEND_WF_NAME = "<temp>.temp_send_workflow"
 
 
+def check_is_in_coroutine() -> bool:
+    try:
+        asyncio.get_running_loop()
+        return True
+    except RuntimeError:
+        return False
+
+
 class WorkflowHandleFuture(Generic[R]):
 
     def __init__(self, workflow_id: str, future: Future[R], dbos: "DBOS"):
@@ -828,6 +837,11 @@ def record_get_result(func: Callable[[], R]) -> R:
             dbos._sys_db.record_get_result(workflow_id, serialized_r, None)
             return r
 
+        if check_is_in_coroutine() and not inspect.iscoroutinefunction(func):
+            dbos_logger.warning(
+                f"Sync workflow ({get_dbos_func_name(func)}) shouldn't be invoked from within another async function. Define it as async or use asyncio.to_thread instead."
+            )
+
         outcome = (
             wfOutcome.wrap(init_wf, dbos=dbos)
             .also(DBOSAssumeRole(rr))
@@ -1009,6 +1023,10 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
                 assert (
                     ctx.is_workflow()
                 ), "Transactions must be called from within workflows"
+                if check_is_in_coroutine():
+                    dbos_logger.warning(
+                        f"Transaction function ({get_dbos_func_name(func)}) shouldn't be invoked from within another async function. Use asyncio.to_thread instead."
+                    )
                 with DBOSAssumeRole(rr):
                     return invoke_tx(*args, **kwargs)
             else:
@@ -1153,6 +1171,10 @@ def check_existing_result() -> Union[NoResult, R]:
 
         @wraps(func)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if check_is_in_coroutine() and not inspect.iscoroutinefunction(func):
+                dbos_logger.warning(
+                    f"Sync step ({get_dbos_func_name(func)}) shouldn't be invoked from within another async function. Define it as async or use asyncio.to_thread instead."
+                )
             # If the step is called from a workflow, run it as a step.
             # Otherwise, run it as a normal function.
             ctx = get_local_dbos_context()

tests/test_async.py

@@ -5,6 +5,9 @@
 
 import pytest
 import sqlalchemy as sa
+from opentelemetry._logs import set_logger_provider
+from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
+from opentelemetry.sdk._logs.export import BatchLogRecordProcessor, InMemoryLogExporter
 
 # Public API
 from dbos import (
@@ -19,6 +22,8 @@
 from dbos._dbos import WorkflowHandle
 from dbos._dbos_config import ConfigFile
 from dbos._error import DBOSAwaitedWorkflowCancelledError, DBOSException
+from dbos._logger import dbos_logger
+from dbos._registrations import get_dbos_func_name
 
 
 @pytest.mark.asyncio
@@ -31,7 +36,7 @@ async def test_async_workflow(dbos: DBOS) -> None:
     async def test_workflow(var1: str, var2: str) -> str:
         nonlocal wf_counter
         wf_counter += 1
-        res1 = test_transaction(var1)
+        res1 = await asyncio.to_thread(test_transaction, var1)
         res2 = await test_step(var2)
         DBOS.logger.info("I'm test_workflow")
         return res1 + res2
@@ -88,7 +93,7 @@ async def test_async_step(dbos: DBOS) -> None:
     async def test_workflow(var1: str, var2: str) -> str:
         nonlocal wf_counter
         wf_counter += 1
-        res1 = test_transaction(var1)
+        res1 = await asyncio.to_thread(test_transaction, var1)
         res2 = await test_step(var2)
         DBOS.logger.info("I'm test_workflow")
         return res1 + res2
@@ -325,6 +330,7 @@ def test_async_tx_raises(config: ConfigFile) -> None:
         async def test_async_tx() -> None:
             pass
 
+    assert "is a coroutine function" in str(exc_info.value)
     # destroy call needed to avoid "functions were registered but DBOS() was not called" warning
     DBOS.destroy(destroy_registry=True)
 
@@ -343,12 +349,12 @@ async def test_workflow(var1: str, var2: str) -> str:
         wf_el_id = id(asyncio.get_running_loop())
         nonlocal wf_counter
         wf_counter += 1
-        res2 = test_step(var2)
+        res2 = await test_step(var2)
         DBOS.logger.info("I'm test_workflow")
         return var1 + res2
 
     @DBOS.step()
-    def test_step(var: str) -> str:
+    async def test_step(var: str) -> str:
         nonlocal step_el_id
         step_el_id = id(asyncio.get_running_loop())
         nonlocal step_counter
@@ -605,3 +611,83 @@ async def run_workflow_task() -> str:
     # Verify the workflow completes despite the task cancellation
     handle: WorkflowHandleAsync[str] = await DBOS.retrieve_workflow_async(wfid)
     assert await handle.get_result() == "completed"
+
+
+@pytest.mark.asyncio
+async def test_check_async_violation(dbos: DBOS) -> None:
+    # Set up in-memory log exporter
+    log_exporter = InMemoryLogExporter()  # type: ignore
+    log_processor = BatchLogRecordProcessor(log_exporter)
+    log_provider = LoggerProvider()
+    log_provider.add_log_record_processor(log_processor)
+    set_logger_provider(log_provider)
+    dbos_logger.addHandler(LoggingHandler(logger_provider=log_provider))
+
+    @DBOS.workflow()
+    def sync_workflow() -> str:
+        return "sync"
+
+    @DBOS.step()
+    def sync_step() -> str:
+        return "step"
+
+    @DBOS.workflow()
+    async def async_workflow_sync_step() -> str:
+        return sync_step()
+
+    @DBOS.transaction()
+    def sync_transaction() -> str:
+        return "txn"
+
+    @DBOS.workflow()
+    async def async_workflow_sync_txn() -> str:
+        return sync_transaction()
+
+    # Call a sync workflow should log a warning
+    sync_workflow()
+
+    log_processor.force_flush(timeout_millis=5000)
+    logs = log_exporter.get_finished_logs()
+    assert len(logs) == 1
+    assert (
+        logs[0].log_record.body is not None
+        and f"Sync workflow ({get_dbos_func_name(sync_workflow)}) shouldn't be invoked from within another async function."
+        in logs[0].log_record.body
+    )
+    log_exporter.clear()
+
+    # Call a sync step from within an async workflow should log a warning
+    await async_workflow_sync_step()
+    log_processor.force_flush(timeout_millis=5000)
+    logs = log_exporter.get_finished_logs()
+    assert len(logs) == 1
+    assert (
+        logs[0].log_record.body is not None
+        and f"Sync step ({get_dbos_func_name(sync_step)}) shouldn't be invoked from within another async function."
+        in logs[0].log_record.body
+    )
+    log_exporter.clear()
+
+    # Directly call a sync step should log a warning
+    sync_step()
+    log_processor.force_flush(timeout_millis=5000)
+    logs = log_exporter.get_finished_logs()
+    assert len(logs) == 1
+    assert (
+        logs[0].log_record.body is not None
+        and f"Sync step ({get_dbos_func_name(sync_step)}) shouldn't be invoked from within another async function."
+        in logs[0].log_record.body
+    )
+    log_exporter.clear()
+
+    # Call a sync transaction from within an async workflow should log a warning
+    await async_workflow_sync_txn()
+    log_processor.force_flush(timeout_millis=5000)
+    logs = log_exporter.get_finished_logs()
+    assert len(logs) == 1
+    assert (
+        logs[0].log_record.body is not None
+        and f"Transaction function ({get_dbos_func_name(sync_transaction)}) shouldn't be invoked from within another async function."
+        in logs[0].log_record.body
+    )
+    log_exporter.clear()

tests/test_async_workflow_management.py

@@ -148,7 +148,7 @@ def step_three(x: int) -> int:
 
     wfid = str(uuid.uuid4())
     with SetWorkflowID(wfid):
-        assert simple_workflow(input_val) == output
+        assert await asyncio.to_thread(simple_workflow, input_val) == output
 
     assert step_one_count == 1
     assert step_two_count == 1

tests/test_dbos.py

@@ -1,5 +1,6 @@
 # mypy: disable-error-code="no-redef"
 
+import asyncio
 import datetime
 import logging
 import os
@@ -1668,17 +1669,17 @@ async def async_step(x: int) -> int:
         assert DBOS.workflow_id is None
         return x
 
-    assert step(5) == 5
+    assert await asyncio.to_thread(step, 5) == 5
     assert await async_step(5) == 5
 
     DBOS(config=config)
 
-    assert step(5) == 5
+    assert await asyncio.to_thread(step, 5) == 5
     assert await async_step(5) == 5
 
     DBOS.launch()
 
-    assert step(5) == 5
+    assert await asyncio.to_thread(step, 5) == 5
     assert await async_step(5) == 5
 
     assert len(DBOS.list_workflows()) == 0

tests/test_workflow_introspection.py

@@ -828,18 +828,18 @@ async def test_callchild_first_asyncio(dbos: DBOS) -> None:
     async def parentWorkflow() -> str:
         handle = await dbos.start_workflow_async(child_workflow)
         child_id = await handle.get_result()
-        stepOne()
-        stepTwo()
+        await stepOne()
+        await stepTwo()
         return child_id
 
     @DBOS.step()
-    def stepOne() -> str:
+    async def stepOne() -> str:
         workflow_id = DBOS.workflow_id
         assert workflow_id is not None
         return workflow_id
 
     @DBOS.step()
-    def stepTwo() -> None:
+    async def stepTwo() -> None:
         return
 
     @DBOS.workflow()