wip: try splitting the functions

Author: pedro-psb

pulpcore/tasking/tasks.py

@@ -73,7 +73,7 @@ def _execute_task(task):
     domain = get_domain()
     try:
         log_task_start(task, domain)
-        task_function = get_task_function(task, in_async_context=False)
+        task_function = get_task_function(task)
         result = task_function()
     except Exception:
         exc_type, exc, tb = sys.exc_info()
@@ -94,7 +94,7 @@ async def _aexecute_task(task):
     await sync_to_async(task.set_running)()
     domain = get_domain()
     try:
-        coroutine = get_task_function(task, in_async_context=True)
+        coroutine = await aget_task_function(task)
         result = await coroutine
     except Exception:
         exc_type, exc, tb = sys.exc_info()
@@ -148,44 +148,56 @@ def log_task_failed(task, exc_type, exc, tb, domain):
     _logger.info("\n".join(traceback.format_list(traceback.extract_tb(tb))))
 
 
-def get_task_function(task, in_async_context):
-    """Wraps task function according to the sync or async context we are in."""
-    module_name, function_name = task.name.rsplit(".", 1)
-    module = importlib.import_module(module_name)
-    func = getattr(module, function_name)
+async def aget_task_function(task):
+    """Get and handle task function running from async context."""
+    func, is_coroutine_fn = _load_function(task.name)
     args = task.enc_args or ()
     kwargs = task.enc_kwargs or {}
     immediate = task.immediate
-    is_coroutine_fn = asyncio.iscoroutinefunction(func)
 
-    if immediate and not is_coroutine_fn:
+    if not immediate:
+        raise ValueError("Non-immediate tasks can't run in async context.")
+    elif not is_coroutine_fn:
         raise ValueError("Immediate tasks must be async functions.")
 
-    if not immediate and in_async_context:
-        raise ValueError("Non-immediate tasks can't run in async context.")
+    coro = func(*args, **kwargs)
+    return _add_timeout_to(coro, task.pk)
 
-    if in_async_context:  # has to be immediate
-        coro = func(*args, **kwargs)
-        coro = asyncio.wait_for(coro, timeout=IMMEDIATE_TIMEOUT)
-        return coro
-    else:
-        if not is_coroutine_fn:
-            return lambda: func(*args, **kwargs)
-
-        # async function in sync context needs adapter
-        async def task_wrapper():  # asyncio.wait_for + async_to_sync requires wrapping
-            coro = func(*args, **kwargs)
-            if immediate:
-                coro = asyncio.wait_for(coro, timeout=IMMEDIATE_TIMEOUT)
-            try:
-                return await coro
-            except asyncio.TimeoutError:
-                msg_template = "Immediate task %s timed out after %s seconds."
-                error_msg = msg_template % (task.pk, IMMEDIATE_TIMEOUT)
-                _logger.info(error_msg)
-                raise RuntimeError(error_msg)
-
-        return async_to_sync(task_wrapper)
+
+def get_task_function(task):
+    """Get and handle task function running from sync context."""
+    func, is_coroutine_fn = _load_function(task.name)
+    args = task.enc_args or ()
+    kwargs = task.enc_kwargs or {}
+    immediate = task.immediate
+
+    # no wrapper required
+    if not is_coroutine_fn:
+        return lambda: func(*args, **kwargs)
+
+    # async function in sync context requires wrapper
+    coro = func(*args, **kwargs)
+    final_coro = _add_timeout_to(coro, task.pk) if immediate else coro
+    return async_to_sync(final_coro)
+
+
+def _load_function(task_name):
+    module_name, function_name = task_name.rsplit(".", 1)
+    module = importlib.import_module(module_name)
+    func = getattr(module, function_name)
+    is_coroutine_fn = asyncio.iscoroutinefunction(func)
+    return func, is_coroutine_fn
+
+
+async def _add_timeout_to(coro, task_pk):
+    coro = asyncio.wait_for(coro, timeout=IMMEDIATE_TIMEOUT)
+    try:
+        return await coro
+    except asyncio.TimeoutError:
+        msg_template = "Immediate task %s timed out after %s seconds."
+        error_msg = msg_template % (task_pk, IMMEDIATE_TIMEOUT)
+        _logger.info(error_msg)
+        raise RuntimeError(error_msg)
 
 
 def dispatch(