【挂起】minor: celery 任务支持apply_async_on_commit

runtime/bamboo-pipeline/pipeline/eri/imp/task.py

@@ -11,16 +11,16 @@
 specific language governing permissions and limitations under the License.
 """
 
-import time
 import logging
+import time
 from typing import Optional
 
 from celery import current_app
-from bamboo_engine.eri.models import ExecuteInterruptPoint, ScheduleInterruptPoint
-
 from pipeline.eri.celery.queues import QueueResolver
-
 from pipeline.eri.models import Process
+from pipeline.eri.utils import apply_async_on_commit
+
+from bamboo_engine.eri.models import ExecuteInterruptPoint, ScheduleInterruptPoint
 
 logger = logging.getLogger("bamboo_engine")
 
@@ -80,7 +80,8 @@ def execute(
         )
 
         def action():
-            result = current_app.tasks[task_name].apply_async(
+            apply_async_on_commit(
+                celery_task=current_app.tasks[task_name],
                 kwargs={
                     "process_id": process_id,
                     "node_id": node_id,
@@ -92,10 +93,9 @@ def action():
                 **route_params,
             )
             logger.info(
-                "[pipeline-trace](root_pipeline: %s) node(%s) execute task %s sended",
+                "[pipeline-trace](root_pipeline: %s) node(%s) execute task sended",
                 root_pipeline_id,
                 node_id,
-                result.id,
             )
 
         _retry_once(action=action)
@@ -126,7 +126,8 @@ def schedule(
         )
 
         def action():
-            result = current_app.tasks[task_name].apply_async(
+            apply_async_on_commit(
+                current_app.tasks[task_name],
                 kwargs={
                     "process_id": process_id,
                     "node_id": node_id,
@@ -137,7 +138,7 @@ def action():
                 },
                 **route_params,
             )
-            logger.info("[pipeline-trace] node(%s) schedule task %s sended", node_id, result.id)
+            logger.info("[pipeline-trace] node(%s) schedule task sended", node_id)
 
         _retry_once(action=action)
 
@@ -170,7 +171,8 @@ def set_next_schedule(
         headers["timestamp"] += schedule_after
 
         def action():
-            result = current_app.tasks[task_name].apply_async(
+            apply_async_on_commit(
+                current_app.tasks[task_name],
                 kwargs={
                     "process_id": process_id,
                     "node_id": node_id,
@@ -182,6 +184,7 @@ def action():
                 countdown=schedule_after,
                 **route_params,
             )
-            logger.info("[pipeline-trace] node(%s) schedule task %s sended", node_id, result.id)
+
+            logger.info("[pipeline-trace] node(%s) schedule task sended", node_id)
 
         _retry_once(action=action)

runtime/bamboo-pipeline/pipeline/eri/utils.py

@@ -10,12 +10,18 @@
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
 """
+import logging
 import socket
-from typing import Dict
+from functools import partial
+from typing import Dict, Optional
 
-from bamboo_engine.eri import ContextValueType
 from celery import current_app
 from django.conf import settings
+from django.db import transaction
+
+from bamboo_engine.eri import ContextValueType
+
+logger = logging.getLogger("bamboo_engine")
 
 CONTEXT_TYPE_MAP = {
     "plain": ContextValueType.PLAIN,
@@ -75,3 +81,29 @@ def check_worker(connection=None):
         tries += 1
 
     return True, worker_list
+
+
+def apply_async_on_commit(celery_task, using: Optional[str] = None, *args, **kwargs):
+    """
+    Apply celery task async and always ignore the task result,
+    it will trigger the task on transaction commit when it is in a atomic block.
+    """
+
+    fn = partial(celery_task.apply_async, ignore_result=True, *args, **kwargs)
+
+    connection = transaction.get_connection(using)
+    if connection.in_atomic_block:
+        logger.debug("trigger task %s on transaction commit", celery_task.name)
+        transaction.on_commit(fn)
+
+    else:
+        logger.debug("trigger task %s immediately", celery_task.name)
+        fn()
+
+
+def delay_on_commit(celery_task, *args, **kwargs):
+    """
+    Star argument version of `apply_async_on_commit`, does not support the extra options.
+    """
+
+    apply_async_on_commit(celery_task, args=args, kwargs=kwargs)

runtime/bamboo-pipeline/pipeline/tests/eri/imp/test_task.py

@@ -11,10 +11,8 @@
 specific language governing permissions and limitations under the License.
 """
 
-from mock import patch, MagicMock
-
 from django.test import TransactionTestCase
-
+from mock import MagicMock, patch
 from pipeline.eri.imp.task import TaskMixin
 from pipeline.eri.models import Process
 
@@ -42,6 +40,7 @@ def test_execute__headers_is_none(self):
                 )
 
         celery_app.tasks["pipeline.eri.celery.tasks.execute"].apply_async.assert_called_once_with(
+            ignore_result=True,
             kwargs={
                 "process_id": 1,
                 "node_id": "nid",
@@ -72,6 +71,7 @@ def test_execute__headers_is_not_none(self):
                 )
 
         celery_app.tasks["pipeline.eri.celery.tasks.execute"].apply_async.assert_called_once_with(
+            ignore_result=True,
             kwargs={
                 "process_id": 1,
                 "node_id": "nid",
@@ -102,6 +102,7 @@ def test_schedule__headers_is_none(self):
                 )
 
         celery_app.tasks["pipeline.eri.celery.tasks.execute"].apply_async.assert_called_once_with(
+            ignore_result=True,
             kwargs={
                 "process_id": 1,
                 "node_id": "nid",
@@ -131,6 +132,7 @@ def test_schedule__headers_is_not_none(self):
                 )
 
         celery_app.tasks["pipeline.eri.celery.tasks.execute"].apply_async.assert_called_once_with(
+            ignore_result=True,
             kwargs={
                 "process_id": 1,
                 "node_id": "nid",
@@ -162,6 +164,7 @@ def test_set_schedule__headers_is_none(self):
                 )
 
         celery_app.tasks["pipeline.eri.celery.tasks.execute"].apply_async.assert_called_once_with(
+            ignore_result=True,
             kwargs={
                 "process_id": 1,
                 "node_id": "nid",
@@ -193,6 +196,7 @@ def test_set_schedule__headers_is_not_none(self):
                 )
 
         celery_app.tasks["pipeline.eri.celery.tasks.execute"].apply_async.assert_called_once_with(
+            ignore_result=True,
             kwargs={
                 "process_id": 1,
                 "node_id": "nid",