feat: 新增流程创建任务并发控制 --story=128208486

app_desc.yaml

@@ -127,7 +127,7 @@ modules:
         plan: 4C4G5R
         replicas: 5
       cworker:
-        command: celery -A blueapps.core.celery worker -P threads -Q celery,pipeline_additional_task,pipeline_additional_task_priority,task_common_${BKFLOW_MODULE_CODE},task_callback_${BKFLOW_MODULE_CODE},node_auto_retry_${BKFLOW_MODULE_CODE},timeout_node_execute_${BKFLOW_MODULE_CODE},timeout_node_record_${BKFLOW_MODULE_CODE} -n common_worker@%h -c 10 -l info
+        command: celery -A blueapps.core.celery worker -P threads -Q celery,pipeline_additional_task,pipeline_additional_task_priority,task_common_${BKFLOW_MODULE_CODE},task_callback_${BKFLOW_MODULE_CODE},task_process_queue_${BKFLOW_MODULE_CODE},node_auto_retry_${BKFLOW_MODULE_CODE},timeout_node_execute_${BKFLOW_MODULE_CODE},timeout_node_record_${BKFLOW_MODULE_CODE} -n common_worker@%h -c 10 -l info
         plan: 4C1G5R
         replicas: 5
       timeout:

bkflow/pipeline_plugins/components/collections/subprocess_plugin/v1_0_0.py

@@ -32,6 +32,7 @@
 from bkflow.contrib.api.collections.interface import InterfaceModuleClient
 from bkflow.exceptions import ValidationError
 from bkflow.pipeline_plugins.components.collections.base import BKFlowBaseService
+from bkflow.task.utils import push_task_to_queue
 
 
 class Subprocess(BaseModel):
@@ -143,7 +144,7 @@ def _create_subprocess_task_instance(self, subprocess, template, pipeline_tree,
         from bkflow.task.utils import extract_extra_info
 
         with transaction.atomic():
-            time_zone = timezone.pytz.timezone(settings.TIME_ZONE) or "Asia/Shanghai"
+            time_zone = timezone.pytz.timezone(settings.TIME_ZONE)
             time_stamp = datetime.datetime.now(tz=time_zone).strftime("%Y%m%d%H%M%S")
             create_task_data = {
                 "name": f"{subprocess.subprocess_name}_子流程_{time_stamp}",
@@ -190,7 +191,7 @@ def _create_subprocess_task_instance(self, subprocess, template, pipeline_tree,
             except TaskFlowRelation.DoesNotExist:
                 root_task_id = parent_task.id
 
-            relate_info = {"node_id": self.id, "node_version": self.version}
+            relate_info = {"node_id": self.id, "node_version": self.version, "parent_task_id": parent_task.id}
             TaskFlowRelation.objects.create(
                 task_id=task_instance.id,
                 parent_task_id=parent_task.id,
@@ -214,6 +215,7 @@ def _create_subprocess_task_instance(self, subprocess, template, pipeline_tree,
     def plugin_execute(self, data, parent_data):
         from bkflow.task.models import TaskInstance
         from bkflow.task.operations import TaskOperation
+        from bkflow.task.utils import task_concurrency_limit_reached
 
         parent_task_id = parent_data.get_one_of_inputs("task_id")
         try:
@@ -235,6 +237,14 @@ def plugin_execute(self, data, parent_data):
 
         # 设置输出并启动任务
         data.set_outputs("task_id", task_instance.id)
+
+        if task_concurrency_limit_reached(task_instance.space_id, task_instance.template_id):
+            try:
+                push_task_to_queue(task_instance, "start")
+            except Exception as e:
+                data.set_outputs("ex_data", str(e))
+                return False
+            return True
         task_operation = TaskOperation(task_instance=task_instance, queue=settings.BKFLOW_MODULE.code)
         operation_method = getattr(task_operation, "start", None)
         if operation_method is None:

bkflow/space/configs.py

@@ -438,6 +438,24 @@ def validate(cls, value: str):
         return True
 
 
+# 流程并发控制
+class ConcurrencyControlConfig(BaseSpaceConfig):
+    name = "concurrency_control"
+    desc = _("流程并发控制")
+    default_value = 0
+    LEAST_NUMBER = 1
+    control = True
+
+    @classmethod
+    def validate(cls, value: str):
+        if int(value) < cls.LEAST_NUMBER:
+            raise ValidationError(
+                f"[validate concurrency control error]: concurrency control only support {cls.LEAST_NUMBER}"
+            )
+
+        return True
+
+
 # 定义 SCHEMA_V1 对应的模型
 class SchemaV1Model(BaseModel):
     meta_apis: str

bkflow/task/celery/tasks.py

@@ -40,7 +40,13 @@
 from bkflow.task.node_timeout import node_timeout_handler
 from bkflow.task.operations import TaskNodeOperation, TaskOperation
 from bkflow.task.serializers import CreateTaskInstanceSerializer
-from bkflow.task.utils import ATOM_FAILED, redis_inst_check, send_task_instance_message
+from bkflow.task.utils import (
+    ATOM_FAILED,
+    push_task_to_queue,
+    redis_inst_check,
+    send_task_instance_message,
+    task_concurrency_limit_reached,
+)
 
 logger = logging.getLogger("celery")
 
@@ -199,6 +205,13 @@ def bkflow_periodic_task_start(*args, **kwargs):
             }
         )
 
+        if task_concurrency_limit_reached(task_instance.space_id, task_instance.template_id):
+            try:
+                push_task_to_queue(task_instance, "start")
+            except Exception as e:
+                logger.exception(f"[bkflow_periodic_task_start] push task to queue failed: {e}")
+            return
+
         task_operation = TaskOperation(task_instance=task_instance, queue=settings.BKFLOW_MODULE.code)
         operation_method = getattr(task_operation, "start")
         if operation_method is None:

bkflow/task/domains/callback.py

@@ -27,16 +27,17 @@
 from pipeline.eri.models import Schedule as DBSchedule
 from pipeline.eri.runtime import BambooDjangoRuntime
 
-from bkflow.task.models import TaskFlowRelation
+from bkflow.task.models import TaskInstance
+from bkflow.task.utils import push_task_to_queue, task_concurrency_limit_reached
 from bkflow.utils.redis_lock import redis_lock
 
 logger = logging.getLogger("root")
 
 
 class TaskCallBacker:
-    def __init__(self, task_id, *args, **kwargs):
+    def __init__(self, task_id, task_relate, *args, **kwargs):
         self.task_id = task_id
-        self.task_relate = TaskFlowRelation.objects.filter(task_id=self.task_id).first()
+        self.task_relate = task_relate
         self.extra_info = {"task_id": self.task_id, **self.task_relate.extra_info, **kwargs}
 
     def check_record_existence(self):
@@ -68,6 +69,19 @@ def subprocess_callback(self):
                     runtime.set_state(node_id=node_id, version=version, to_state=states.READY)
                     runtime.set_state(node_id=node_id, version=version, to_state=states.RUNNING)
 
+                parent_task_id = self.extra_info["parent_task_id"]
+                parent_task = TaskInstance.objects.filter(id=parent_task_id).first()
+
+                if (
+                    task_concurrency_limit_reached(parent_task.space_id, parent_task.template_id, is_exemption=True)
+                    and self.extra_info["task_success"] is True
+                ):
+                    try:
+                        push_task_to_queue(parent_task, "callback")
+                    except Exception as e:
+                        logger.exception(f"[TaskCallBacker _subprocess_callback] push task to queue error: {e}")
+                    return True
+
                 bamboo_engine_api.callback(runtime=runtime, node_id=node_id, version=version, data=self.extra_info)
 
         except Exception as e:

bkflow/task/models.py

@@ -330,8 +330,7 @@ def change_parent_task_node_state_to_running(self):
             runtime.set_execution_data_outputs(parent_node_id, data_outputs)
 
             # 仅当父流程的节点状态为失败时，才需要唤醒父流程的节点
-            parent_task_id = TaskFlowRelation.objects.filter(task_id=self.id).first().parent_task_id
-            parent_task = TaskInstance.objects.get(id=parent_task_id)
+            parent_task = TaskInstance.objects.get(id=record.parent_task_id)
             parent_task.change_parent_task_node_state_to_running()
 
 

bkflow/task/serializers.py

@@ -118,10 +118,11 @@ class TaskInstanceSerializer(serializers.ModelSerializer):
     create_time = serializers.DateTimeField(format="%Y-%m-%d %H:%M:%S%z")
     start_time = serializers.DateTimeField(format="%Y-%m-%d %H:%M:%S%z")
     finish_time = serializers.DateTimeField(format="%Y-%m-%d %H:%M:%S%z")
+    is_waiting = serializers.SerializerMethodField()
 
     class Meta:
         model = TaskInstance
-        fields = "__all__"
+        exclude = ["extra_info"]
         read_only_fields = (
             "id",
             "instance_id",
@@ -141,6 +142,9 @@ class Meta:
             "tree_info_id",
         )
 
+    def get_is_waiting(self, instance):
+        return instance.extra_info.get("is_waiting", False)
+
 
 class RetrieveTaskInstanceSerializer(TaskInstanceSerializer):
     pipeline_tree = serializers.SerializerMethodField()

bkflow/task/signals/handlers.py

@@ -33,7 +33,12 @@
     TaskInstance,
     TimeoutNodeConfig,
 )
-from bkflow.task.utils import ATOM_FAILED, TASK_FINISHED, redis_inst_check
+from bkflow.task.utils import (
+    ATOM_FAILED,
+    TASK_FINISHED,
+    process_task_from_queue,
+    redis_inst_check,
+)
 
 logger = logging.getLogger("root")
 
@@ -99,12 +104,14 @@ def bamboo_engine_eri_post_set_state_handler(sender, node_id, to_state, version,
             queue=f"task_common_{settings.BKFLOW_MODULE.code}",
             routing_key=f"task_common_{settings.BKFLOW_MODULE.code}",
         )
+        _process_task_from_queue(root_id)
     elif to_state == bamboo_engine_states.REVOKED and node_id == root_id:
         try:
             TaskInstance.objects.set_revoked(root_id)
         except Exception as e:
             logger.exception(f"TaskInstance set revoked error: {e}")
         _check_and_callback(root_id, task_success=False)
+        _process_task_from_queue(root_id)
     elif to_state == bamboo_engine_states.FINISHED and node_id == root_id:
         try:
             TaskInstance.objects.set_finished(root_id)
@@ -119,13 +126,28 @@ def bamboo_engine_eri_post_set_state_handler(sender, node_id, to_state, version,
             routing_key=f"task_common_{settings.BKFLOW_MODULE.code}",
         )
         _check_and_callback(root_id, task_success=True)
+        _process_task_from_queue(root_id)
 
     try:
         _node_timeout_info_update(settings.redis_inst, to_state, node_id, version)
     except Exception as e:
         logger.exception(f"node_timeout_info_update error: {e}")
 
 
+def _process_task_from_queue(root_id):
+    try:
+        template_id = TaskInstance.objects.get(instance_id=root_id).template_id
+        process_task_from_queue.apply_async(
+            kwargs={
+                "template_id": template_id,
+            },
+            queue=f"task_process_queue_{settings.BKFLOW_MODULE.code}",
+            routing_key=f"task_process_queue_{settings.BKFLOW_MODULE.code}",
+        )
+    except Exception as e:
+        logger.exception(f"TaskInstance get template_id error: {e}")
+
+
 def _check_and_callback(instance_id, *args, **kwargs):
     try:
         task_id = TaskInstance.objects.get(instance_id=instance_id).id
@@ -143,7 +165,7 @@ def task_callback(task_id, retry_times=0, *args, **kwargs):
     task_relate = TaskFlowRelation.objects.filter(task_id=task_id).first()
     if not task_relate:
         return
-    tcb = TaskCallBacker(task_id, *args, **kwargs)
+    tcb = TaskCallBacker(task_id, task_relate, *args, **kwargs)
     if not tcb.check_record_existence():
         message = f"[task_callback] task_id {task_id} does not in TaskCallBackRecord."
         logger.error(message)

bkflow/task/utils.py

@@ -22,6 +22,7 @@
 from functools import wraps
 
 from bamboo_engine import states as bamboo_engine_states
+from celery import current_app
 from django.conf import settings
 from django.utils.translation import ugettext_lazy as _
 from pipeline.core import constants as pipeline_constants
@@ -30,6 +31,8 @@
 
 from bkflow.utils.dates import format_datetime
 from bkflow.utils.message import send_message
+from bkflow.utils.space import space_config_manager
+from bkflow.utils.trace import start_trace
 
 logger = logging.getLogger("root")
 
@@ -166,3 +169,126 @@ def extract_extra_info(constants, keys=None):
     for key in list(constants.keys()) if not keys else keys:
         extra_info.update({key: {"name": constants[key]["name"], "value": constants[key]["value"]}})
     return json.dumps(extra_info, ensure_ascii=False)
+
+
+@redis_inst_check
+def push_task_to_queue(task, operation, node_id=None, data=None):
+    template_id = task.template_id
+    redis_key = f"task_wait_{template_id}"
+
+    # 准备任务数据
+    task_data = {"operation": operation, "task_id": task.id}
+    if node_id:
+        task_data.update({"node_id": node_id})
+    if data:
+        task_data.update({"node_data": data})
+    task_json = json.dumps(task_data)
+
+    lua_script = """
+        local queue_key = KEYS[1]
+        local max_size = tonumber(ARGV[1])
+        local task_data = ARGV[2]
+
+        local current_size = redis.call('llen', queue_key)
+        if current_size >= max_size then
+            return -1  -- 队列已满
+        end
+
+        redis.call('rpush', queue_key, task_data)
+        return current_size + 1  -- 返回新队列大小
+    """
+
+    with start_trace("push_task_to_queue", operation=operation, task_id=task.id):
+        result = settings.redis_inst.eval(lua_script, 1, redis_key, settings.TASK_QUEUE_MAX_SIZE, task_json)
+
+        if result == -1:
+            logger.error(f"Task queue for template {template_id} is full, cannot add more tasks")
+            raise Exception(f"Task queue for template {template_id} is full, cannot add more tasks")
+
+        logger.info(f"Task {task.id} added to queue for template {template_id}, new queue size: {result}")
+
+    task.extra_info.update({"is_waiting": True})
+    task.save()
+    return True
+
+
+@current_app.task()
+@redis_inst_check
+def process_task_from_queue(template_id):
+    from bkflow.task.models import TaskInstance
+    from bkflow.task.operations import TaskNodeOperation, TaskOperation
+
+    redis_key = f"task_wait_{template_id}"
+    task_json = settings.redis_inst.lpop(redis_key)
+    if not task_json:
+        return None
+
+    task_data = json.loads(task_json)
+    operation = task_data.get("operation")
+    task_instance = TaskInstance.objects.get(id=task_data.get("task_id"))
+
+    for invoke_num in range(1, settings.TASK_MAX_RETRY_FREQUENCY + 1):
+        task_instance.extra_info.update({"is_waiting": False})
+        try:
+            if operation in ["start", "resume"]:
+                task_operation = TaskOperation(task_instance, settings.BKFLOW_MODULE.code)
+                operation_method = getattr(task_operation, operation, None)
+            else:
+                node_operation = TaskNodeOperation(task_instance, task_data.get("node_id"))
+                operation_method = getattr(node_operation, operation, None)
+
+            operation_result = operation_method(operator=operation, **task_data.get("node_data", {}))
+            if operation_result.result:
+                break
+            opera_error = operation_result.message
+        except Exception as e:
+            logger.error(f"Failed to process task {task_instance.id} from queue (attempt {invoke_num}): {e}")
+            opera_error = e
+
+        if invoke_num == settings.TASK_MAX_RETRY_FREQUENCY and opera_error:
+            logger.error(f"Failed to process task {task_instance.id} kwargs {task_data} from queue")
+            task_instance.extra_info.update({"operation_failed": opera_error})
+
+    task_instance.save()
+    return task_instance
+
+
+def get_running_task_count(space_id, template_id):
+    """统计当前正在执行的任务数量"""
+    from bkflow.task.models import TaskInstance
+    from bkflow.task.operations import TaskOperation
+
+    task_instances = TaskInstance.objects.filter(
+        space_id=space_id, template_id=template_id, is_deleted=False, is_started=True, is_finished=False
+    )
+
+    task_operations = [
+        {"task_id": task_instance.id, "operation": TaskOperation(task_instance=task_instance).get_task_states()}
+        for task_instance in task_instances
+    ]
+
+    task_count = 0
+    for task_operation in task_operations:
+        if task_operation["operation"].result is False:
+            continue
+        if task_operation["operation"].data.get("state") == "RUNNING":
+            task_count += 1
+
+    return task_count
+
+
+def task_concurrency_limit_reached(space_id, template_id, is_exemption=False):
+    """判断是否超出并发限制"""
+    concurrency_control = space_config_manager.get_concurrency_control(space_id)
+    if not concurrency_control:
+        return False
+
+    redis_key = f"task_wait_{template_id}"
+    queue_size = settings.redis_inst.llen(redis_key)
+    if queue_size != 0:
+        return True
+
+    running_count = get_running_task_count(space_id, template_id)
+    if is_exemption:
+        return running_count > concurrency_control
+    return running_count >= concurrency_control