diff --git a/runtime/bamboo-pipeline/pipeline/eri/imp/task.py b/runtime/bamboo-pipeline/pipeline/eri/imp/task.py index 86bb385f..a97a2dea 100644 --- a/runtime/bamboo-pipeline/pipeline/eri/imp/task.py +++ b/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) diff --git a/runtime/bamboo-pipeline/pipeline/eri/utils.py b/runtime/bamboo-pipeline/pipeline/eri/utils.py index 687e803c..fd9c8b53 100644 --- a/runtime/bamboo-pipeline/pipeline/eri/utils.py +++ b/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) diff --git a/runtime/bamboo-pipeline/pipeline/tests/eri/imp/test_task.py b/runtime/bamboo-pipeline/pipeline/tests/eri/imp/test_task.py index 3956f7d9..21111457 100644 --- a/runtime/bamboo-pipeline/pipeline/tests/eri/imp/test_task.py +++ b/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",