diff --git a/.env.example b/.env.example index 562a17c8..72cfa33e 100644 --- a/.env.example +++ b/.env.example @@ -105,12 +105,12 @@ WINDUP_MQ_MAX_CONSUME_ATTEMPTS=5 WINDUP_MQ_CONSUME_LEASE_SECONDS=1800 WINDUP_MQ_EMAIL_HANDLER_RETRIES=3 # worker 内 handler 并行度（单进程内 ThreadPoolExecutor） -# 生成期间不占 Postgres 连接，可按机器内存与上游配额上调。 +# 图 / 动作分两条 Stream 独立出队。默认对齐现网，避免无环境变量时打满上游。 WINDUP_MQ_EMAIL_CONCURRENCY=8 -WINDUP_MQ_GENERATION_IMAGE_CONCURRENCY=16 -WINDUP_MQ_GENERATION_ACTION_CONCURRENCY=8 -# poll 使用独立线程池,不与图/动作共享槽位 -WINDUP_MQ_GENERATION_POLL_CONCURRENCY=16 +WINDUP_MQ_GENERATION_IMAGE_CONCURRENCY=4 +WINDUP_MQ_GENERATION_ACTION_CONCURRENCY=2 +# poll 使用独立线程池,不与图/动作提交共享槽位 +WINDUP_MQ_GENERATION_POLL_CONCURRENCY=2 # i2v 轮询走 ZSET 延迟队列,到期促进到 Stream;不用 Redis 过期事件 WINDUP_MQ_DELAYED_ZSET=windup:zset:delayed WINDUP_MQ_DELAYED_CLAIM_LIMIT=50 diff --git a/backend/packages/app/src/windup_app/bootstrap/worker.py b/backend/packages/app/src/windup_app/bootstrap/worker.py index 6e4024e6..578050b1 100644 --- a/backend/packages/app/src/windup_app/bootstrap/worker.py +++ b/backend/packages/app/src/windup_app/bootstrap/worker.py @@ -7,7 +7,13 @@ import threading import time -from windup_app.server.mq.catalog import all_stream_specs, email_stream_spec, generation_stream_spec +from windup_app.server.mq.catalog import ( + all_stream_specs, + email_stream_spec, + generation_action_stream_spec, + generation_image_stream_spec, + generation_stream_spec, +) from windup_app.server.orchestrator import task_repo from windup_app.server.orchestrator.executor import ( bind_matte, @@ -67,25 +73,31 @@ def _handle_signal(_signum, _frame) -> None: signal.signal(signal.SIGTERM, _handle_signal) signal.signal(signal.SIGINT, _handle_signal) - consumers = [ - StreamConsumer( - email_stream_spec(), + def _generation_consumer(spec): + return StreamConsumer( + spec, run_image_task=run_image_task, run_action_task=run_action_task, run_direction_set_task=run_direction_set_task, run_view_sheet_task=run_view_sheet_task, stop_event=stop_event, - ), + resume_action_poll=resume_action_poll, + resume_action_client_bake=resume_action_client_bake, + ) + + consumers = [ StreamConsumer( - generation_stream_spec(), + email_stream_spec(), run_image_task=run_image_task, run_action_task=run_action_task, run_direction_set_task=run_direction_set_task, run_view_sheet_task=run_view_sheet_task, stop_event=stop_event, - resume_action_poll=resume_action_poll, - resume_action_client_bake=resume_action_client_bake, ), + _generation_consumer(generation_image_stream_spec()), + _generation_consumer(generation_action_stream_spec()), + # 过渡 drain：切流前已进旧 generation Stream 的消息还要被消费。 + _generation_consumer(generation_stream_spec()), ] threads = [consumer.start() for consumer in consumers] relay_thread = start_relay_loop(stop_event) @@ -99,7 +111,10 @@ def _handle_signal(_signum, _frame) -> None: ) pending_thread.start() - logger.info("windup worker 已启动 | streams=%s", [s.stream for s in all_stream_specs()]) + logger.info( + "windup worker 已启动 | streams=%s", + [s.stream for s in (*all_stream_specs(), generation_stream_spec())], + ) try: while not stop_event.is_set(): diff --git a/backend/packages/app/src/windup_app/server/mq/catalog.py b/backend/packages/app/src/windup_app/server/mq/catalog.py index 783808ee..d3b7ef4c 100644 --- a/backend/packages/app/src/windup_app/server/mq/catalog.py +++ b/backend/packages/app/src/windup_app/server/mq/catalog.py @@ -14,6 +14,7 @@ - ``recover_as``：PENDING 任务按此 GenerationType 值重入队；轮询类留 None 3. 在 handlers 的 ``HANDLERS`` 登记可调用对象 +图片与动作分两条 Stream，出队互不堵。旧 ``GENERATION_STREAM`` 只给过渡 drain。 不在此表里塞积分账本或 SSE EventBus。是否新开 Stream 仍按 SLA 决定。 """ @@ -31,9 +32,13 @@ EMAIL_STREAM = "windup:stream:email" GENERATION_STREAM = "windup:stream:generation" +GENERATION_IMAGE_STREAM = "windup:stream:generation-image" +GENERATION_ACTION_STREAM = "windup:stream:generation-action" EMAIL_GROUP = "email" GENERATION_GROUP = "generation" +GENERATION_IMAGE_GROUP = "generation-image" +GENERATION_ACTION_GROUP = "generation-action" MSG_TYPE_VERIFICATION_CODE = "verification_code" MSG_TYPE_CHARACTER_IMAGE = "character_image" @@ -72,17 +77,17 @@ class TypeSpec: def generation_image_concurrency() -> int: - # executor 已短 session:生成/上传不再占连接,默认不再被 15 连接池卡住。 - return _env_int("WINDUP_MQ_GENERATION_IMAGE_CONCURRENCY", 16) + # 默认对齐现网：图与动作分队后各自限流，避免无环境变量时按 16/8 打上游。 + return _env_int("WINDUP_MQ_GENERATION_IMAGE_CONCURRENCY", 4) def generation_action_concurrency() -> int: - return _env_int("WINDUP_MQ_GENERATION_ACTION_CONCURRENCY", 8) + return _env_int("WINDUP_MQ_GENERATION_ACTION_CONCURRENCY", 2) def generation_poll_concurrency() -> int: - # 单次 inspect + 偶尔下载,不 sleep,默认高于 action 建单并发。 - return _env_int("WINDUP_MQ_GENERATION_POLL_CONCURRENCY", 16) + # 单次 inspect + 偶尔下载,不 sleep；与动作提交分池，不占提交槽。 + return _env_int("WINDUP_MQ_GENERATION_POLL_CONCURRENCY", 2) def type_specs() -> tuple[TypeSpec, ...]: @@ -96,7 +101,7 @@ def type_specs() -> tuple[TypeSpec, ...]: ), TypeSpec( msg_type=MSG_TYPE_CHARACTER_IMAGE, - stream=GENERATION_STREAM, + stream=GENERATION_IMAGE_STREAM, pool=POOL_SHARED, concurrency=generation_image_concurrency(), limit=True, @@ -104,7 +109,7 @@ def type_specs() -> tuple[TypeSpec, ...]: ), TypeSpec( msg_type=MSG_TYPE_CHARACTER_ACTION, - stream=GENERATION_STREAM, + stream=GENERATION_ACTION_STREAM, pool=POOL_SHARED, concurrency=generation_action_concurrency(), limit=True, @@ -112,14 +117,14 @@ def type_specs() -> tuple[TypeSpec, ...]: ), TypeSpec( msg_type=MSG_TYPE_CHARACTER_ACTION_POLL, - stream=GENERATION_STREAM, + stream=GENERATION_ACTION_STREAM, pool=POOL_POLL, concurrency=generation_poll_concurrency(), limit=True, ), TypeSpec( msg_type=MSG_TYPE_CHARACTER_ACTION_CLIENT_BAKE, - stream=GENERATION_STREAM, + stream=GENERATION_ACTION_STREAM, pool=POOL_POLL, concurrency=generation_poll_concurrency(), limit=True, @@ -135,7 +140,23 @@ def type_spec(msg_type: str) -> TypeSpec | None: def types_for_stream(stream: str) -> tuple[TypeSpec, ...]: - return tuple(spec for spec in type_specs() if spec.stream == stream) + matched = tuple(spec for spec in type_specs() if spec.stream == stream) + if stream != GENERATION_STREAM: + return matched + # 旧产线 Stream 过渡 drain：切流后仍要按图/动作类型处理残留消息。 + drained = tuple( + spec + for spec in type_specs() + if spec.stream in (GENERATION_IMAGE_STREAM, GENERATION_ACTION_STREAM) + ) + return matched + drained + + +def stream_for_msg_type(msg_type: str) -> str: + spec = type_spec(msg_type) + if spec is None: + raise ValueError(f"未知消息类型: {msg_type}") + return spec.stream def msg_type_for_generation(task_type: str) -> str: @@ -157,8 +178,8 @@ def msg_type_for_generation(task_type: str) -> str: def _pool_size(stream: str, pool: str) -> int: return sum( spec.concurrency - for spec in type_specs() - if spec.stream == stream and spec.pool == pool + for spec in types_for_stream(stream) + if spec.pool == pool ) @@ -170,13 +191,31 @@ def email_stream_spec() -> StreamSpec: ) +def generation_image_stream_spec() -> StreamSpec: + return StreamSpec( + stream=GENERATION_IMAGE_STREAM, + group=GENERATION_IMAGE_GROUP, + concurrency=_pool_size(GENERATION_IMAGE_STREAM, POOL_SHARED), + ) + + +def generation_action_stream_spec() -> StreamSpec: + return StreamSpec( + stream=GENERATION_ACTION_STREAM, + group=GENERATION_ACTION_GROUP, + concurrency=_pool_size(GENERATION_ACTION_STREAM, POOL_SHARED), + ) + + def generation_worker_pool_size() -> int: - # image/action 共用一个线程池。poll 走独立 pool 名,不能加进这个数字, - # 否则 image 占满线程后 poll 只能排队。 - return _pool_size(GENERATION_STREAM, POOL_SHARED) + # 旧 Stream drain 的共享池：图+动作合计。poll 不进这个数字。 + return _pool_size(GENERATION_IMAGE_STREAM, POOL_SHARED) + _pool_size( + GENERATION_ACTION_STREAM, POOL_SHARED + ) def generation_stream_spec() -> StreamSpec: + """旧 ``generation`` Stream 的过渡 drain 规格。新任务不要再往这里投。""" return StreamSpec( stream=GENERATION_STREAM, group=GENERATION_GROUP, @@ -184,8 +223,12 @@ def generation_stream_spec() -> StreamSpec: ) -def all_stream_specs() -> tuple[StreamSpec, StreamSpec]: - return email_stream_spec(), generation_stream_spec() +def all_stream_specs() -> tuple[StreamSpec, ...]: + return ( + email_stream_spec(), + generation_image_stream_spec(), + generation_action_stream_spec(), + ) __all__ = [ @@ -194,9 +237,14 @@ def all_stream_specs() -> tuple[StreamSpec, StreamSpec]: "EMAIL_HANDLER_RETRIES", "GENERATION_STREAM", "GENERATION_GROUP", + "GENERATION_IMAGE_STREAM", + "GENERATION_IMAGE_GROUP", + "GENERATION_ACTION_STREAM", + "GENERATION_ACTION_GROUP", "GENERATION_PENDING_MAX_AGE_SECONDS", "GENERATION_RUNNING_STALE_SECONDS", "MSG_TYPE_CHARACTER_ACTION", + "MSG_TYPE_CHARACTER_ACTION_CLIENT_BAKE", "MSG_TYPE_CHARACTER_ACTION_POLL", "MSG_TYPE_CHARACTER_IMAGE", "MSG_TYPE_VERIFICATION_CODE", @@ -208,11 +256,14 @@ def all_stream_specs() -> tuple[StreamSpec, StreamSpec]: "all_stream_specs", "email_stream_spec", "generation_action_concurrency", + "generation_action_stream_spec", "generation_image_concurrency", + "generation_image_stream_spec", "generation_poll_concurrency", "generation_stream_spec", "generation_worker_pool_size", "msg_type_for_generation", + "stream_for_msg_type", "type_spec", "type_specs", "types_for_stream", diff --git a/backend/packages/app/src/windup_app/server/orchestrator/client_bake.py b/backend/packages/app/src/windup_app/server/orchestrator/client_bake.py index 19f2b5db..88dbb0ab 100644 --- a/backend/packages/app/src/windup_app/server/orchestrator/client_bake.py +++ b/backend/packages/app/src/windup_app/server/orchestrator/client_bake.py @@ -19,8 +19,8 @@ from dataclasses import asdict, dataclass from windup_app.server.mq.catalog import ( - GENERATION_STREAM, MSG_TYPE_CHARACTER_ACTION_CLIENT_BAKE, + stream_for_msg_type, ) from windup_framework.db.redis import get_redis from windup_framework.mq.delayed import schedule_delayed @@ -93,7 +93,7 @@ def open_job(task_id: int, spec: ClientBakeSpec) -> float: pipe.execute() schedule_delayed( delay_s=DEADLINE_S, - stream=GENERATION_STREAM, + stream=stream_for_msg_type(MSG_TYPE_CHARACTER_ACTION_CLIENT_BAKE), msg_type=MSG_TYPE_CHARACTER_ACTION_CLIENT_BAKE, payload={"task_id": task_id, "reason": REASON_TIMEOUT}, dedupe_key=f"generation:{task_id}:clientbake:timeout", @@ -169,7 +169,7 @@ def schedule_resume(task_id: int, reason: str = REASON_FRAMES, detail: str = "") payload["detail"] = detail[:200] schedule_delayed( delay_s=0, - stream=GENERATION_STREAM, + stream=stream_for_msg_type(MSG_TYPE_CHARACTER_ACTION_CLIENT_BAKE), msg_type=MSG_TYPE_CHARACTER_ACTION_CLIENT_BAKE, payload=payload, dedupe_key=f"generation:{task_id}:clientbake:{reason}", diff --git a/backend/packages/app/src/windup_app/server/orchestrator/i2v_poll.py b/backend/packages/app/src/windup_app/server/orchestrator/i2v_poll.py index a1e2d225..03385072 100644 --- a/backend/packages/app/src/windup_app/server/orchestrator/i2v_poll.py +++ b/backend/packages/app/src/windup_app/server/orchestrator/i2v_poll.py @@ -12,7 +12,7 @@ from dataclasses import dataclass from typing import Any -from windup_app.server.mq.catalog import GENERATION_STREAM, MSG_TYPE_CHARACTER_ACTION_POLL +from windup_app.server.mq.catalog import MSG_TYPE_CHARACTER_ACTION_POLL, stream_for_msg_type from windup_app.server.mq.i2v_state import ( I2V_FIRST_POLL_S, I2V_MAX_WAIT_S, @@ -91,7 +91,7 @@ def schedule( ) schedule_delayed( delay_s=wait, - stream=GENERATION_STREAM, + stream=stream_for_msg_type(MSG_TYPE_CHARACTER_ACTION_POLL), msg_type=MSG_TYPE_CHARACTER_ACTION_POLL, payload=_poll_payload(task_id, poll_count), dedupe_key=_poll_dedupe(task_id, poll_count), @@ -149,7 +149,7 @@ def reschedule_if_waiting(task_id: int, *, delay_s: float = 1) -> bool: poll_count = int(state.get("poll_count") or 0) schedule_delayed( delay_s=delay_s, - stream=GENERATION_STREAM, + stream=stream_for_msg_type(MSG_TYPE_CHARACTER_ACTION_POLL), msg_type=MSG_TYPE_CHARACTER_ACTION_POLL, payload=_poll_payload(task_id, poll_count), dedupe_key=_poll_dedupe(task_id, poll_count), diff --git a/backend/packages/app/src/windup_app/server/orchestrator/recover.py b/backend/packages/app/src/windup_app/server/orchestrator/recover.py index c78a70e1..86bade52 100644 --- a/backend/packages/app/src/windup_app/server/orchestrator/recover.py +++ b/backend/packages/app/src/windup_app/server/orchestrator/recover.py @@ -16,8 +16,8 @@ from windup_app.server.mq.catalog import ( GENERATION_RUNNING_STALE_SECONDS, - GENERATION_STREAM, msg_type_for_generation, + stream_for_msg_type, ) from windup_app.server.orchestrator import billing, task_repo from windup_app.server.orchestrator.i2v_poll import reschedule_if_waiting @@ -109,10 +109,11 @@ def _requeue_pending( if hasattr(task.task_type, "value") else str(task.task_type) ) + msg_type = msg_type_for_generation(task_type) message_id = publisher.enqueue( session, - stream=GENERATION_STREAM, - msg_type=msg_type_for_generation(task_type), + stream=stream_for_msg_type(msg_type), + msg_type=msg_type, payload={ "task_id": task.id, "task_type": task_type, diff --git a/backend/packages/app/src/windup_app/web/api/generation.py b/backend/packages/app/src/windup_app/web/api/generation.py index 76637ef6..edc0b14a 100644 --- a/backend/packages/app/src/windup_app/web/api/generation.py +++ b/backend/packages/app/src/windup_app/web/api/generation.py @@ -44,8 +44,8 @@ from windup_app.server.character.model import Character, CharacterData from windup_app.server.orchestrator import billing, task_repo from windup_app.server.mq.catalog import ( - GENERATION_STREAM, msg_type_for_generation, + stream_for_msg_type, ) from windup_app.server.orchestrator.service import service as generation_service from windup_app.server.orchestrator.model import ( @@ -531,7 +531,7 @@ def _publish_generation_after_commit( msg_type = msg_type_for_generation(task_type) message_id = publisher.enqueue( session, - stream=GENERATION_STREAM, + stream=stream_for_msg_type(msg_type), msg_type=msg_type, payload={"task_id": task_id, "task_type": task_type}, dedupe_key=dedupe_key or f"generation:{task_id}", diff --git a/backend/packages/app/src/windup_app/worker/consumer.py b/backend/packages/app/src/windup_app/worker/consumer.py index 157d80c7..9c965ad7 100644 --- a/backend/packages/app/src/windup_app/worker/consumer.py +++ b/backend/packages/app/src/windup_app/worker/consumer.py @@ -96,6 +96,9 @@ def __init__( self._image_sem = self._semaphores.get(MSG_TYPE_CHARACTER_IMAGE) self._action_sem = self._semaphores.get(MSG_TYPE_CHARACTER_ACTION) self._poll_sem = self._semaphores.get(MSG_TYPE_CHARACTER_ACTION_POLL) + self._pool_capacity = {name: max(1, size) for name, size in pool_sizes.items()} + self._inflight_by_pool = {name: 0 for name in self._pool_capacity} + self._inflight_lock = threading.Lock() self._claim_cursor = "0-0" self._last_claim_at = 0.0 @@ -127,6 +130,10 @@ def _loop(self) -> None: self._claim_idle(redis_client) self._last_claim_at = now + if not self._has_free_slot(): + self._stop.wait(timeout=0.05) + continue + try: batches = mq_client.xreadgroup( redis_client, @@ -142,25 +149,90 @@ def _loop(self) -> None: for _stream, messages in batches: for stream_id, fields in messages: - self._submit_message(stream_id, fields) + self._accept_message(stream_id, fields) + + def _has_free_slot(self) -> bool: + """任一线程池还有空槽才继续 XREADGROUP。具体类型在领到之后再卡对应池。""" + with self._inflight_lock: + return any( + self._inflight_by_pool[pool] < self._pool_capacity[pool] + for pool in self._pool_capacity + ) + + def _pool_has_slot(self, pool: str) -> bool: + cap = self._pool_capacity.get(pool) + if cap is None: + cap = self._pool_capacity[POOL_SHARED] + pool = POOL_SHARED + with self._inflight_lock: + return self._inflight_by_pool.get(pool, 0) < cap + + def _acquire_slot(self, pool: str) -> None: + if pool not in self._pool_capacity: + pool = POOL_SHARED + with self._inflight_lock: + self._inflight_by_pool[pool] = self._inflight_by_pool.get(pool, 0) + 1 + + def _release_slot(self, pool: str) -> None: + if pool not in self._pool_capacity: + pool = POOL_SHARED + with self._inflight_lock: + n = self._inflight_by_pool.get(pool, 0) + if n > 0: + self._inflight_by_pool[pool] = n - 1 + + def _pool_name_for(self, fields: dict[str, str]) -> str: + try: + msg_type = str(mq_client.parse_envelope(fields)["type"]) + except Exception: + return POOL_SHARED + spec = type_spec(msg_type) + if spec is None or spec.pool not in self._pool_capacity: + return POOL_SHARED + return spec.pool + + def _accept_message(self, stream_id: str, fields: dict[str, str]) -> None: + """已领取的消息按目标池等空槽再提交，避免 PEL 里堆着执行器排队。""" + pool = self._pool_name_for(fields) + while not self._pool_has_slot(pool) and not self._stop.is_set(): + self._stop.wait(timeout=0.05) + if self._stop.is_set(): + return + self._submit_message(stream_id, fields) def _claim_idle(self, redis_client) -> None: while not self._stop.is_set(): + if not self._has_free_slot(): + break claimed, next_start = mq_client.claim_idle_messages( redis_client, self._config.stream, self._config.group, self._consumer_name, start_id=self._claim_cursor, + count=1, ) self._claim_cursor = next_start for stream_id, fields in claimed: - self._submit_message(stream_id, fields) + self._accept_message(stream_id, fields) if not claimed: break def _submit_message(self, stream_id: str, fields: dict[str, str]) -> None: - self._executor_for(fields).submit(self._process_message, stream_id, fields) + pool = self._pool_name_for(fields) + self._acquire_slot(pool) + try: + future = self._executor_for(fields).submit( + self._process_message, stream_id, fields + ) + except Exception: + self._release_slot(pool) + raise + done = getattr(future, "add_done_callback", None) + if done is None: + self._release_slot(pool) + return + done(lambda _f: self._release_slot(pool)) def _executor_for(self, fields: dict[str, str]) -> ThreadPoolExecutor: try: diff --git a/backend/tests/test_generation_api.py b/backend/tests/test_generation_api.py index 0b5cd8c0..6d89210b 100644 --- a/backend/tests/test_generation_api.py +++ b/backend/tests/test_generation_api.py @@ -371,6 +371,7 @@ def test_four_view_copies_confirmed_master_and_shares_image_queue(auth_client): assert "direction" not in body["data"]["input_payload"] assert publisher.enqueue.call_args.kwargs["msg_type"] == "character_image" assert publisher.enqueue.call_args.kwargs["payload"]["task_type"] == "character_four_view" + assert publisher.enqueue.call_args.kwargs["stream"] == "windup:stream:generation-image" def test_four_view_without_confirmed_master_is_rejected_before_queueing(auth_client): diff --git a/backend/tests/test_generation_quota.py b/backend/tests/test_generation_quota.py index 497d00f5..cbacd133 100644 --- a/backend/tests/test_generation_quota.py +++ b/backend/tests/test_generation_quota.py @@ -424,6 +424,7 @@ def test_recover_requeues_pending_tasks_with_open_freeze(session_factory): assert len(enqueued) == 1 assert enqueued[0]["dedupe_key"] == f"generation:{task_id}" assert enqueued[0]["msg_type"] == "character_action" + assert enqueued[0]["stream"] == "windup:stream:generation-action" assert enqueued[0]["payload"]["task_id"] == task_id with session_factory() as session: @@ -616,6 +617,7 @@ def test_recover_requeues_pending_image_tasks(session_factory): assert len(enqueued) == 1 assert enqueued[0]["msg_type"] == "character_image" + assert enqueued[0]["stream"] == "windup:stream:generation-image" assert enqueued[0]["payload"]["task_id"] == task_id diff --git a/backend/tests/test_i2v_poll.py b/backend/tests/test_i2v_poll.py index 576b219e..9bf7e8a6 100644 --- a/backend/tests/test_i2v_poll.py +++ b/backend/tests/test_i2v_poll.py @@ -33,6 +33,7 @@ def test_schedule_persists_state_and_enqueues_delayed_poll(monkeypatch): assert saved["task_id"] == 9 assert saved["job_id"] == "j1" assert delayed["msg_type"] == "character_action_poll" + assert delayed["stream"] == "windup:stream:generation-action" assert delayed["payload"]["task_id"] == 9 assert delayed["dedupe_key"] == "generation:9:poll:0" diff --git a/backend/tests/test_mq_catalog.py b/backend/tests/test_mq_catalog.py index 4e2dc639..d7cefbcc 100644 --- a/backend/tests/test_mq_catalog.py +++ b/backend/tests/test_mq_catalog.py @@ -7,9 +7,14 @@ from windup_app.server.mq.catalog import ( EMAIL_GROUP, EMAIL_STREAM, + GENERATION_ACTION_GROUP, + GENERATION_ACTION_STREAM, GENERATION_GROUP, + GENERATION_IMAGE_GROUP, + GENERATION_IMAGE_STREAM, GENERATION_STREAM, MSG_TYPE_CHARACTER_ACTION, + MSG_TYPE_CHARACTER_ACTION_CLIENT_BAKE, MSG_TYPE_CHARACTER_ACTION_POLL, MSG_TYPE_CHARACTER_IMAGE, MSG_TYPE_VERIFICATION_CODE, @@ -18,13 +23,17 @@ all_stream_specs, email_stream_spec, generation_action_concurrency, + generation_action_stream_spec, generation_image_concurrency, + generation_image_stream_spec, generation_poll_concurrency, generation_stream_spec, generation_worker_pool_size, msg_type_for_generation, + stream_for_msg_type, type_spec, type_specs, + types_for_stream, ) @@ -35,31 +44,40 @@ def test_email_stream_spec_defaults(): assert spec.concurrency == 8 -def test_generation_stream_spec_aggregates_pool_size(): - spec = generation_stream_spec() - assert spec.stream == GENERATION_STREAM - assert spec.group == GENERATION_GROUP - assert spec.concurrency == generation_worker_pool_size() - assert spec.concurrency == ( +def test_generation_streams_are_split_and_drain_keeps_old_key(): + image = generation_image_stream_spec() + action = generation_action_stream_spec() + drain = generation_stream_spec() + assert image.stream == GENERATION_IMAGE_STREAM + assert image.group == GENERATION_IMAGE_GROUP + assert image.concurrency == generation_image_concurrency() + assert action.stream == GENERATION_ACTION_STREAM + assert action.group == GENERATION_ACTION_GROUP + assert action.concurrency == generation_action_concurrency() + assert drain.stream == GENERATION_STREAM + assert drain.group == GENERATION_GROUP + assert drain.concurrency == generation_worker_pool_size() + assert drain.concurrency == ( generation_image_concurrency() + generation_action_concurrency() ) -def test_default_generation_concurrency_is_the_documented_floor(monkeypatch): - """短 session 后默认不再被 15 连接卡住;16 图 + 8 动作是面向多用户的底线。""" +def test_default_generation_concurrency_matches_production_floor(monkeypatch): monkeypatch.delenv("WINDUP_MQ_GENERATION_IMAGE_CONCURRENCY", raising=False) monkeypatch.delenv("WINDUP_MQ_GENERATION_ACTION_CONCURRENCY", raising=False) monkeypatch.delenv("WINDUP_MQ_GENERATION_POLL_CONCURRENCY", raising=False) - assert generation_image_concurrency() == 16 - assert generation_action_concurrency() == 8 - assert generation_poll_concurrency() == 16 + assert generation_image_concurrency() == 4 + assert generation_action_concurrency() == 2 + assert generation_poll_concurrency() == 2 -def test_all_stream_specs_returns_email_and_generation(): +def test_all_stream_specs_returns_email_image_and_action(): specs = all_stream_specs() - assert len(specs) == 2 - assert specs[0].stream == EMAIL_STREAM - assert specs[1].stream == GENERATION_STREAM + assert [spec.stream for spec in specs] == [ + EMAIL_STREAM, + GENERATION_IMAGE_STREAM, + GENERATION_ACTION_STREAM, + ] def test_catalog_respects_env_overrides(monkeypatch): @@ -72,23 +90,50 @@ def test_catalog_respects_env_overrides(monkeypatch): assert generation_image_concurrency() == 10 assert generation_action_concurrency() == 5 assert generation_poll_concurrency() == 4 + assert generation_image_stream_spec().concurrency == 10 + assert generation_action_stream_spec().concurrency == 5 assert generation_stream_spec().concurrency == 15 -def test_type_specs_register_pool_limit_and_recover(): +def test_type_specs_register_pool_limit_recover_and_stream(): by_type = {spec.msg_type: spec for spec in type_specs()} assert by_type[MSG_TYPE_VERIFICATION_CODE].stream == EMAIL_STREAM assert by_type[MSG_TYPE_VERIFICATION_CODE].pool == POOL_SHARED assert by_type[MSG_TYPE_VERIFICATION_CODE].limit is False + assert by_type[MSG_TYPE_CHARACTER_IMAGE].stream == GENERATION_IMAGE_STREAM assert by_type[MSG_TYPE_CHARACTER_IMAGE].pool == POOL_SHARED assert by_type[MSG_TYPE_CHARACTER_IMAGE].limit is True assert by_type[MSG_TYPE_CHARACTER_IMAGE].recover_as == MSG_TYPE_CHARACTER_IMAGE + assert by_type[MSG_TYPE_CHARACTER_ACTION].stream == GENERATION_ACTION_STREAM assert by_type[MSG_TYPE_CHARACTER_ACTION].recover_as == MSG_TYPE_CHARACTER_ACTION + assert by_type[MSG_TYPE_CHARACTER_ACTION_POLL].stream == GENERATION_ACTION_STREAM assert by_type[MSG_TYPE_CHARACTER_ACTION_POLL].pool == POOL_POLL assert by_type[MSG_TYPE_CHARACTER_ACTION_POLL].recover_as is None + assert by_type[MSG_TYPE_CHARACTER_ACTION_CLIENT_BAKE].stream == GENERATION_ACTION_STREAM assert type_spec("unknown") is None +def test_stream_for_msg_type_routes_image_and_action(): + assert stream_for_msg_type(MSG_TYPE_CHARACTER_IMAGE) == GENERATION_IMAGE_STREAM + assert stream_for_msg_type(MSG_TYPE_CHARACTER_ACTION) == GENERATION_ACTION_STREAM + assert stream_for_msg_type(MSG_TYPE_CHARACTER_ACTION_POLL) == GENERATION_ACTION_STREAM + assert stream_for_msg_type(MSG_TYPE_CHARACTER_ACTION_CLIENT_BAKE) == GENERATION_ACTION_STREAM + assert stream_for_msg_type(MSG_TYPE_VERIFICATION_CODE) == EMAIL_STREAM + with pytest.raises(ValueError, match="未知消息类型"): + stream_for_msg_type("unknown") + + +def test_legacy_generation_stream_still_sees_image_and_action_types(): + names = {spec.msg_type for spec in types_for_stream(GENERATION_STREAM)} + assert names >= { + MSG_TYPE_CHARACTER_IMAGE, + MSG_TYPE_CHARACTER_ACTION, + MSG_TYPE_CHARACTER_ACTION_POLL, + MSG_TYPE_CHARACTER_ACTION_CLIENT_BAKE, + } + assert MSG_TYPE_VERIFICATION_CODE not in names + + def test_msg_type_for_generation_skips_poll_types(): assert msg_type_for_generation(MSG_TYPE_CHARACTER_IMAGE) == MSG_TYPE_CHARACTER_IMAGE assert msg_type_for_generation("character_direction_set") == MSG_TYPE_CHARACTER_IMAGE diff --git a/backend/tests/test_mq_worker.py b/backend/tests/test_mq_worker.py index ec58e4e6..86ef8a5b 100644 --- a/backend/tests/test_mq_worker.py +++ b/backend/tests/test_mq_worker.py @@ -13,10 +13,14 @@ from conftest import seed_credit_account from windup_app.server.mq.catalog import ( + GENERATION_ACTION_STREAM, + GENERATION_IMAGE_STREAM, MSG_TYPE_CHARACTER_ACTION, MSG_TYPE_CHARACTER_ACTION_POLL, MSG_TYPE_CHARACTER_IMAGE, MSG_TYPE_VERIFICATION_CODE, + POOL_POLL, + POOL_SHARED, ) from windup_app.server.orchestrator import billing, task_repo from windup_app.server.orchestrator.model import ( @@ -541,6 +545,7 @@ def submit(self, fn, stream_id, fields): def fake_claim(*_args, **_kwargs): claim_calls["count"] += 1 + assert _kwargs.get("count") == 1 if claim_calls["count"] == 1: return ([("2-0", {"data": "{}"})], "2-0") return ([], "2-0") @@ -608,6 +613,163 @@ def fake_xreadgroup(*_args, **_kwargs): thread.join(timeout=3) +def test_consumer_skips_xreadgroup_when_image_slots_full(engine, monkeypatch): + monkeypatch.setenv("WINDUP_MQ_GENERATION_IMAGE_CONCURRENCY", "1") + _patch_worker_session_local(monkeypatch, engine) + + stop = threading.Event() + redis_mock = MagicMock() + monkeypatch.setattr("windup_app.worker.consumer.get_redis", lambda: redis_mock) + monkeypatch.setattr( + "windup_app.worker.consumer.mq_client.ensure_consumer_group", + lambda *_a: None, + ) + monkeypatch.setattr( + "windup_app.worker.consumer.mq_client.claim_idle_messages", + lambda *_a, **_k: ([], "0-0"), + ) + xread_calls: list[int] = [] + + def fake_xreadgroup(*_args, **_kwargs): + xread_calls.append(1) + stop.wait(timeout=0.2) + return [] + + monkeypatch.setattr("windup_app.worker.consumer.mq_client.xreadgroup", fake_xreadgroup) + + started = threading.Event() + release = threading.Event() + + def fake_process(_self, _stream_id, _fields): + started.set() + release.wait(timeout=5) + + monkeypatch.setattr(StreamConsumer, "_process_message", fake_process) + + consumer = StreamConsumer( + ConsumerConfig(stream=GENERATION_IMAGE_STREAM, group="generation-image", concurrency=1), + run_image_task=MagicMock(), + run_action_task=MagicMock(), + stop_event=stop, + ) + try: + consumer._submit_message( + "i-0", + { + "data": json.dumps( + { + "v": 1, + "id": str(uuid.uuid4()), + "type": MSG_TYPE_CHARACTER_IMAGE, + "payload": {"task_id": 1, "task_type": "character_image"}, + } + ) + }, + ) + assert started.wait(timeout=5) + assert consumer._has_free_slot() is False + thread = consumer.start() + assert stop.wait(timeout=0.3) is False + assert xread_calls == [] + release.set() + stop.set() + thread.join(timeout=3) + finally: + release.set() + stop.set() + consumer.shutdown() + + +def test_consumer_does_not_queue_action_beyond_shared_pool(engine, monkeypatch): + """动作共享池满时，多领到的 action 不得进执行器队列；poll 槽空着也不能拿去堆 action。""" + monkeypatch.setenv("WINDUP_MQ_GENERATION_ACTION_CONCURRENCY", "1") + monkeypatch.setenv("WINDUP_MQ_GENERATION_POLL_CONCURRENCY", "2") + _patch_worker_session_local(monkeypatch, engine) + + stop = threading.Event() + redis_mock = MagicMock() + monkeypatch.setattr("windup_app.worker.consumer.get_redis", lambda: redis_mock) + monkeypatch.setattr( + "windup_app.worker.consumer.mq_client.ensure_consumer_group", + lambda *_a: None, + ) + monkeypatch.setattr( + "windup_app.worker.consumer.mq_client.claim_idle_messages", + lambda *_a, **_k: ([], "0-0"), + ) + + first_started = threading.Event() + first_release = threading.Event() + second_started = threading.Event() + + def fake_process(_self, stream_id, _fields): + if stream_id == "a-0": + first_started.set() + first_release.wait(timeout=5) + return + second_started.set() + stop.set() + + monkeypatch.setattr(StreamConsumer, "_process_message", fake_process) + + extra = { + "data": json.dumps( + { + "v": 1, + "id": str(uuid.uuid4()), + "type": MSG_TYPE_CHARACTER_ACTION, + "payload": {"task_id": 2, "task_type": "character_action"}, + } + ) + } + + def fake_xreadgroup(*_args, **_kwargs): + if stop.is_set() or second_started.is_set(): + return [] + return [(GENERATION_ACTION_STREAM, [("a-1", extra)])] + + monkeypatch.setattr("windup_app.worker.consumer.mq_client.xreadgroup", fake_xreadgroup) + + consumer = StreamConsumer( + ConsumerConfig( + stream=GENERATION_ACTION_STREAM, + group="generation-action", + concurrency=1, + ), + run_image_task=MagicMock(), + run_action_task=MagicMock(), + stop_event=stop, + ) + try: + consumer._submit_message( + "a-0", + { + "data": json.dumps( + { + "v": 1, + "id": str(uuid.uuid4()), + "type": MSG_TYPE_CHARACTER_ACTION, + "payload": {"task_id": 1, "task_type": "character_action"}, + } + ) + }, + ) + assert first_started.wait(timeout=5) + assert consumer._pool_has_slot(POOL_SHARED) is False + assert consumer._pool_has_slot(POOL_POLL) is True + assert consumer._has_free_slot() is True + thread = consumer.start() + assert second_started.wait(timeout=0.4) is False + first_release.set() + assert second_started.wait(timeout=5) + stop.set() + thread.join(timeout=3) + finally: + first_release.set() + stop.set() + consumer.shutdown() + + def test_start_relay_loop_invokes_relay(monkeypatch): relay_calls: list[int] = [] stop = threading.Event()