From 7a4601c94450c8459317c15d8e36b787925df37a Mon Sep 17 00:00:00 2001 From: guohelu <19503896967@163.com> Date: Wed, 27 May 2026 17:23:33 +0800 Subject: [PATCH 1/2] =?UTF-8?q?feat:=20=E6=B5=81=E7=A8=8B=E4=BB=BB?= =?UTF-8?q?=E5=8A=A1=E6=94=AF=E6=8C=81=E5=B9=B6=E5=8F=91=E9=99=90=E5=88=B6?= =?UTF-8?q?=20--story=3D134711116=20#=20Reviewed,=20transaction=20id:=2081?= =?UTF-8?q?044?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../collections/subprocess_plugin/v1_0_0.py | 9 ++- bkflow/space/configs.py | 17 ++++ bkflow/task/celery/tasks.py | 7 ++ bkflow/task/signals/handlers.py | 21 ++++- bkflow/task/utils.py | 77 +++++++++++++++++++ bkflow/task/views.py | 21 +++++ 6 files changed, 149 insertions(+), 3 deletions(-) diff --git a/bkflow/pipeline_plugins/components/collections/subprocess_plugin/v1_0_0.py b/bkflow/pipeline_plugins/components/collections/subprocess_plugin/v1_0_0.py index 5a98eac001..238bbbed28 100644 --- a/bkflow/pipeline_plugins/components/collections/subprocess_plugin/v1_0_0.py +++ b/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 diff --git a/bkflow/space/configs.py b/bkflow/space/configs.py index 5bb42e4ac2..a3660a234c 100644 --- a/bkflow/space/configs.py +++ b/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 diff --git a/bkflow/task/celery/tasks.py b/bkflow/task/celery/tasks.py index c94b4245cc..b0922d88e1 100644 --- a/bkflow/task/celery/tasks.py +++ b/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() diff --git a/bkflow/task/signals/handlers.py b/bkflow/task/signals/handlers.py index d532b87638..3e93205525 100644 --- a/bkflow/task/signals/handlers.py +++ b/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: diff --git a/bkflow/task/utils.py b/bkflow/task/utils.py index 8606b08343..f55ae53c81 100644 --- a/bkflow/task/utils.py +++ b/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, "" diff --git a/bkflow/task/views.py b/bkflow/task/views.py index 8428fa88af..cf50238e60 100644 --- a/bkflow/task/views.py +++ b/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="任务状态查询") From 33852dd6f59607d2551be90cbb6511499cce9fe8 Mon Sep 17 00:00:00 2001 From: guohelu <19503896967@163.com> Date: Wed, 27 May 2026 17:55:46 +0800 Subject: [PATCH 2/2] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E5=8D=95=E4=BE=A7?= =?UTF-8?q?=E5=A4=B1=E8=B4=A5=E9=97=AE=E9=A2=98=20--story=3D134711116=20#?= =?UTF-8?q?=20Reviewed,=20transaction=20id:=2081049?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/interface/space/test_space_config.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/interface/space/test_space_config.py b/tests/interface/space/test_space_config.py index 52a8cf8f29..5c79dcd533 100644 --- a/tests/interface/space/test_space_config.py +++ b/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