fix(scheduling): 负载均衡模式与无亲和键场景统一走随机排序

重构 CandidateSorter.shuffle_keys_by_internal_priority 中同优先级
Key 的排序逻辑,将三分支简化为两分支:
- 随机排序:TTL=0 / 负载均衡模式 / 无 affinity_key
- 哈希确定性排序:缓存亲和模式且有 affinity_key

移除了原先"无 affinity_key 时按 ID 排序"的冗余分支,新增
对应单元测试覆盖三种场景。
This commit is contained in:
fawney19
2026-03-02 22:36:24 +08:00
parent 8e98eed5c8
commit e26ed8481f
2 changed files with 104 additions and 9 deletions

View File

@@ -0,0 +1,93 @@
from __future__ import annotations
from types import SimpleNamespace
from typing import cast
from unittest.mock import patch
from src.models.database import ProviderAPIKey
from src.services.scheduling.candidate_sorter import CandidateSorter
from src.services.scheduling.scheduling_config import SchedulingConfig
from src.services.scheduling.utils import affinity_hash
def _make_key(key_id: str, internal_priority: int = 1) -> ProviderAPIKey:
return cast(
ProviderAPIKey,
SimpleNamespace(
id=key_id,
internal_priority=internal_priority,
),
)
def _reverse_in_place(items: list[object]) -> None:
items.reverse()
def test_shuffle_keys_random_in_load_balance_mode_even_with_affinity_key() -> None:
config = SchedulingConfig(
priority_mode=SchedulingConfig.PRIORITY_MODE_PROVIDER,
scheduling_mode=SchedulingConfig.SCHEDULING_MODE_LOAD_BALANCE,
)
sorter = CandidateSorter(config)
keys = [_make_key("k1"), _make_key("k2"), _make_key("k3")]
with patch(
"src.services.scheduling.candidate_sorter.random.shuffle",
side_effect=_reverse_in_place,
) as shuffle_mock:
result = sorter.shuffle_keys_by_internal_priority(
keys,
affinity_key="affinity-1",
use_random=False,
)
assert [k.id for k in result] == ["k3", "k2", "k1"]
assert shuffle_mock.call_count == 1
def test_shuffle_keys_random_when_affinity_key_absent() -> None:
config = SchedulingConfig(
priority_mode=SchedulingConfig.PRIORITY_MODE_PROVIDER,
scheduling_mode=SchedulingConfig.SCHEDULING_MODE_CACHE_AFFINITY,
)
sorter = CandidateSorter(config)
keys = [_make_key("k1"), _make_key("k2"), _make_key("k3")]
with patch(
"src.services.scheduling.candidate_sorter.random.shuffle",
side_effect=_reverse_in_place,
) as shuffle_mock:
result = sorter.shuffle_keys_by_internal_priority(
keys,
affinity_key=None,
use_random=False,
)
assert [k.id for k in result] == ["k3", "k2", "k1"]
assert shuffle_mock.call_count == 1
def test_shuffle_keys_still_hashes_with_affinity_in_non_load_balance_mode() -> None:
config = SchedulingConfig(
priority_mode=SchedulingConfig.PRIORITY_MODE_PROVIDER,
scheduling_mode=SchedulingConfig.SCHEDULING_MODE_CACHE_AFFINITY,
)
sorter = CandidateSorter(config)
keys = [_make_key("k1"), _make_key("k2"), _make_key("k3")]
affinity_key = "affinity-1"
expected = sorted(keys, key=lambda k: affinity_hash(affinity_key, k.id))
with patch("src.services.scheduling.candidate_sorter.random.shuffle") as shuffle_mock:
result = sorter.shuffle_keys_by_internal_priority(
keys,
affinity_key=affinity_key,
use_random=False,
)
assert [k.id for k in result] == [k.id for k in expected]
shuffle_mock.assert_not_called()