feat: 流程任务支持并发限制 --story=134711116

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

@@ -33,6 +33,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 check_template_concurrency, update_running_task
 from bkflow.utils.handlers import mask_sensitive_data_for_display
 
 
@@ -300,7 +301,13 @@ def plugin_execute(self, data, parent_data):
         operation_method = getattr(task_operation, "start", None)
         if operation_method is None:
             raise ValidationError("task operation not found")
-        operation_method(operator=parent_task.creator)
+        is_allowed, msg = check_template_concurrency(task_instance.space_id, task_instance.template_id)
+        if not is_allowed:
+            data.set_outputs("ex_data", msg)
+            return False
+        operation_result = operation_method(operator=parent_task.creator)
+        if operation_result.result:
+            update_running_task(task_instance.template_id, task_instance.id, action="add")
 
         return True
 

bkflow/space/configs.py

@@ -440,6 +440,23 @@ 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

@@ -43,8 +43,10 @@
 from bkflow.task.utils import (
     ATOM_FAILED,
     add_node_name_to_status_tree,
+    check_template_concurrency,
     redis_inst_check,
     send_task_instance_message,
+    update_running_task,
 )
 from bkflow.utils.json import safe_for_json
 
@@ -268,9 +270,14 @@ def bkflow_periodic_task_start(*args, **kwargs):
         operation_method = getattr(task_operation, "start")
         if operation_method is None:
             raise ValidationError("task operation not found")
+        is_allowed, msg = check_template_concurrency(task_instance.space_id, task_instance.template_id)
+        if not is_allowed:
+            logger.error(f"[bkflow_periodic_task_start] error: {msg}")
+            raise ValidationError(f"[bkflow_periodic_task_start] error: {msg}")
         result = operation_method(operator=periodic_task.creator)
         if result.result:
             logger.info(f"[bkflow_periodic_task_start] task {task_instance.id} started")
+            update_running_task(task_instance.template_id, task_instance.id, action="add")
             periodic_task.total_run_count += 1
             periodic_task.last_run_at = timezone.localtime(timezone.now())
             periodic_task.save()

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,
+    redis_inst_check,
+    update_running_task,
+)
 
 logger = logging.getLogger("root")
 
@@ -84,6 +89,16 @@ def _node_timeout_info_update(redis_inst, to_state, node_id, version):
         redis_inst.zrem(settings.EXECUTING_NODE_POOL, key)
 
 
+def _remove_task_from_running_set(root_id):
+    """根据 pipeline_instance_id (root_id) 查出任务实例，并从 Redis 运行中任务集合中移除"""
+    try:
+        task_instance = TaskInstance.objects.filter(instance_id=root_id).only("id", "template_id").first()
+        if task_instance and task_instance.template_id is not None:
+            update_running_task(task_instance.template_id, task_instance.id, action="remove")
+    except Exception as e:
+        logger.exception(f"[_remove_task_from_running_set] root_id={root_id} error: {e}")
+
+
 @receiver(post_set_state)
 def bamboo_engine_eri_post_set_state_handler(sender, node_id, to_state, version, root_id, parent_id, loop, **kwargs):
     if to_state == bamboo_engine_states.FAILED:
@@ -100,12 +115,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}",
         )
+        _remove_task_from_running_set(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)
+        _remove_task_from_running_set(root_id)
     elif to_state == bamboo_engine_states.FINISHED and node_id == root_id:
         try:
             TaskInstance.objects.set_finished(root_id)
@@ -121,7 +138,7 @@ 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)
-
+        _remove_task_from_running_set(root_id)
     try:
         _node_timeout_info_update(settings.redis_inst, to_state, node_id, version)
     except Exception as e:

bkflow/task/utils.py

@@ -28,6 +28,7 @@
 from pipeline.engine.utils import calculate_elapsed_time
 from redis.client import Redis
 
+from bkflow.contrib.api.collections.interface import InterfaceModuleClient
 from bkflow.utils.dates import format_datetime
 from bkflow.utils.message import send_message
 
@@ -166,3 +167,79 @@ 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)
+
+
+def get_running_tasks_redis_key(template_id) -> str:
+    return f"bkflow_running_tasks_template_{template_id}"
+
+
+@redis_inst_check
+def get_running_task_count(template_id) -> int:
+    """获取指定模板当前运行中的任务数量。Redis 不可用时返回 0，不阻塞流程。"""
+    if template_id is None:
+        return 0
+    try:
+        return int(settings.redis_inst.scard(get_running_tasks_redis_key(template_id)) or 0)
+    except Exception as e:
+        logger.exception(f"[get_running_task_count] template_id={template_id} error: {e}")
+        return 0
+
+
+@redis_inst_check
+def update_running_task(template_id, task_id, action: str) -> bool:
+    """更新模板运行中任务集合。
+
+    :param template_id: 模板 ID
+    :param task_id: 任务 ID
+    :param action: 操作类型，"add"：加入运行集合；"remove"：从运行集合移除
+    :return: 操作是否成功
+    """
+    if template_id is None:
+        return False
+    if action not in ("add", "remove"):
+        logger.error(f"[update_running_task] invalid action: {action}")
+        return False
+    try:
+        key = get_running_tasks_redis_key(template_id)
+        if action == "add":
+            settings.redis_inst.sadd(key, str(task_id))
+        else:
+            settings.redis_inst.srem(key, str(task_id))
+        return True
+    except Exception as e:
+        logger.exception(
+            f"[update_running_task] action={action} template_id={template_id} task_id={task_id} error: {e}"
+        )
+        return False
+
+
+def check_template_concurrency(space_id, template_id):
+    """检查模板当前运行中任务数是否达到空间配置的最大并发上限。
+
+    :param space_id: 空间 ID
+    :param template_id: 模板 ID
+    :return: (is_allowed, message)
+        is_allowed: True 表示允许执行；False 表示已达上限。
+        message: 当 is_allowed 为 False 时为提示信息。
+    """
+    if template_id is None:
+        return True, ""
+    try:
+        interface_client = InterfaceModuleClient()
+        space_infos_result = interface_client.get_space_infos(
+            {"space_id": space_id, "config_names": "concurrency_control"}
+        )
+        space_configs = space_infos_result.get("data", {}).get("configs", {})
+        max_running = int(space_configs.get("concurrency_control", 0))
+    except Exception as e:
+        logger.exception(f"[check_template_concurrency] get space concurrency_control error: {e}")
+        return True, ""
+
+    if not max_running:
+        return True, ""
+
+    running_count = get_running_task_count(template_id)
+    if running_count >= max_running:
+        msg = f"当前模板(template_id={template_id})正在运行的任务数已达上限({max_running})，" f"当前运行中数量为{running_count}，请稍后再试"
+        return False, msg
+    return True, ""

bkflow/task/views.py

@@ -75,6 +75,7 @@
     TaskUpdateLabelSerializer,
     UpdatePeriodicTaskSerializer,
 )
+from bkflow.task.utils import check_template_concurrency, update_running_task
 from bkflow.utils.handlers import handle_plain_log
 from bkflow.utils.mixins import BKFLOWCommonMixin
 from bkflow.utils.permissions import AdminPermission, AppInternalPermission
@@ -276,6 +277,11 @@ def operate(self, request, operation, *args, **kwargs):
             operation_method = getattr(task_operation, operation, None)
             if operation_method is None:
                 raise ValidationError("task operation not found")
+            if operation == "start":
+                is_allowed, msg = check_template_concurrency(task_instance.space_id, task_instance.template_id)
+                if not is_allowed:
+                    return Response({"result": False, "message": msg, "data": None})
+
             data = request.data
             operator = data.pop("operator", request.user.username)
             operation_result = operation_method(operator=operator, **data)
@@ -293,6 +299,11 @@ def operate(self, request, operation, *args, **kwargs):
                         },
                     }
                 )
+            if operation_result.result and operation == "start":
+                update_running_task(task_instance.template_id, task_instance.id, action="add")
+            elif operation_result.result and operation == "revoke":
+                update_running_task(task_instance.template_id, task_instance.id, action="remove")
+
             return Response(dict(operation_result))
 
     @swagger_auto_schema(methods=["post"], operation_description="节点操作", request_body=EmptyBodySerializer)
@@ -320,7 +331,17 @@ def node_operate(self, request, node_id, operation, *args, **kwargs):
                 raise ValidationError("node operation not found")
             data = request.data
             operator = data.pop("operator", request.user.username)
+
+            if operation in ["retry", "skip"] and task_instance.template_id is not None:
+                is_allowed, msg = check_template_concurrency(task_instance.space_id, task_instance.template_id)
+                if not is_allowed:
+                    return Response({"result": False, "message": msg, "data": None})
+
             operation_result = operation_method(operator=operator, **data)
+
+            if operation in ["retry", "skip"] and operation_result.result:
+                update_running_task(task_instance.template_id, task_instance.id, action="add")
+
             return Response(dict(operation_result))
 
     @swagger_auto_schema(methods=["get"], operation_description="任务状态查询")

tests/interface/space/test_space_config.py

@@ -47,9 +47,9 @@
 class TestSpaceConfigHandler:
     def test_get_all_configs(self):
         configs = SpaceConfigHandler.get_all_configs()
-        assert len(configs) == 12
+        assert len(configs) == 13
         configs = SpaceConfigHandler.get_all_configs(only_public=True)
-        assert len(configs) == 11
+        assert len(configs) == 12
 
     def test_get_config(self):
         # valid cases